diff --git a/re2oapi/client.py b/re2oapi/client.py index 88f2c08..624ae20 100644 --- a/re2oapi/client.py +++ b/re2oapi/client.py @@ -2,6 +2,9 @@ import requests from requests.exceptions import HTTPError import datetime import iso8601 +from pathlib import Path +import json +import stat from . import endpoints from . import exceptions @@ -11,19 +14,68 @@ TIME_FOR_RENEW = 60 class Re2oAPIClient: - def __init__(self, hostname, username, password, use_tls=True): + def __init__(self, hostname, username, password, token_file=None, use_tls=True): self.use_tls = use_tls + self.token_file = token_file or Path.home() / '.re2o.token' self.hostname = hostname self._username = username self._password = password - self.token = self._obtain_token() + try: + self.token = self._get_token_from_file() + except Exception: + self.token = self._get_token_from_server() + + def __del__(self): + try: + self._save_token_to_file() + except Exception: + pass - def _need_renew_token(self): + @property + def need_renew_token(self): return self.token['expiration'] > \ datetime.datetime.now(datetime.timezone.utc) + \ datetime.timedelta(seconds=TIME_FOR_RENEW) - def _obtain_token(self): + def _get_token_from_file(self): + if not self.token_file.is_file(): + raise exceptions.TokenFileNotFound(str(self.token_file)) + + with self.token_file.open() as f: + data = json.load(f) + + try: + token_data = data[self.hostname][self._username] + return { + 'token': token_data['token'], + 'expiration': iso8601.parse_date(token_data['expiration']) + } + except KeyError: + raise exceptions.TokenNotInTokenFile( + "{user}@{host} in {token_file}".format(user=self._username, + hostname=self.hostname, + token_file=str(self.token_file)) + ) + + def _save_token_to_file(self): + try: + with self.token_file.open() as f: + data = json.load(f) + except Exception: + data = {} + + if self.hostname not in data.keys(): + data[self.hostname] = {} + data[self.hostname][self._username] = { + 'token': self.token['token'], + 'expiration': self.token['expiration'].isoformat() + } + + with self.token_file.open('w') as f: + json.dump(data, f) + self.token_file.chmod(stat.S_IWRITE | stat.S_IREAD) + + def _get_token_from_server(self): response = requests.post( self.get_url_for('token', forma=None), data={'username': self._username, 'password': self._password} @@ -38,8 +90,8 @@ class Re2oAPIClient: } def get_token(self): - if self._need_renew_token(): - self.token = self._obtain_token() + if self.need_renew_token: + self.token = self._get_token_from_server() return self.token['token'] def get(self, url, *args, **kwargs): diff --git a/re2oapi/exceptions.py b/re2oapi/exceptions.py index 6ff579b..ac76966 100644 --- a/re2oapi/exceptions.py +++ b/re2oapi/exceptions.py @@ -6,3 +6,9 @@ class URLParameterMissing(ValueError): class InvalidCredentials(ValueError): pass + +class TokenFileNotFound(ValueError): + pass + +class TokenNotInTokenFile(ValueError): + pass