Save token in file

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):

re2oapi/exceptions.py

@@ -6,3 +6,9 @@ class URLParameterMissing(ValueError):
 
 class InvalidCredentials(ValueError):
     pass
+
+class TokenFileNotFound(ValueError):
+    pass
+
+class TokenNotInTokenFile(ValueError):
+    pass