Updates to messaging and shadows to use new authentication (#94)

* updates to messaging and shadows to use new authentication

* fixed webbrowser requirement

* Update __init__.py

* Moved authenticate to its own module to avoid circular imports

* fixed auth header

---------

Co-authored-by: Kateryna Voitiuk <[email protected]>

.gitignore

@@ -172,3 +172,6 @@ tmp/
 **/.DS_Store
 dist/
 **/.vscode/**
+
+# don't commit the service_account file
+src/braingeneers/iot/service_account/config.json

pyproject.toml

@@ -51,7 +51,7 @@ dependencies = [
     "typing_extensions>=4.6; python_version<'3.11'",
     'diskcache',
     'pytz', 
-    'tzlocal'
+    'tzlocal',
 ]
 
 [tool.hatch.build.hooks.vcs]

src/braingeneers/iot/__init__.py

@@ -1,4 +1,4 @@
-import braingeneers
-from braingeneers.iot.messaging import *
-from braingeneers.iot.device import *
-from braingeneers.iot.simple import *
+import braingeneers
+from braingeneers.iot.messaging import *
+from braingeneers.iot.device import *
+from braingeneers.iot.simple import *

src/braingeneers/iot/authenticate.py

@@ -0,0 +1,130 @@
+
+import os
+import json
+import webbrowser
+import importlib.resources
+import configparser
+import datetime
+import requests
+import argparse
+
+
+def authenticate_and_get_token():
+    """
+    Directs users to a URL to authenticate and get a JWT token.
+    Once the token has been obtained manually it will refresh automatically every month.
+    By default, the token is valid for 4 months from issuance.
+    Returns token data as a dict containing `access_token` and `expires_at` keys.
+    """
+    PACKAGE_NAME = "braingeneers.iot"
+
+    url = 'https://service-accounts.braingeneers.gi.ucsc.edu/generate_token'
+    print(f'Please visit the following URL to generate your JWT token: {url}')
+    webbrowser.open(url)
+
+    token_json = input('Please paste the JSON token issued by the page and press Enter:\n')
+    try:
+        token_data = json.loads(token_json)
+    except json.JSONDecodeError:
+        raise ValueError('Invalid JSON. Please make sure you have copied the token correctly.')
+
+    config_dir = os.path.join(importlib.resources.files(PACKAGE_NAME), 'service_account')
+    os.makedirs(config_dir, exist_ok=True)
+    config_file = os.path.join(config_dir, 'config.json')
+
+    with open(config_file, 'w') as f:
+        json.dump(token_data, f)
+
+    print('Token has been saved successfully.')
+    return token_data
+
+
+def update_config_file(file_path, section, key, new_value):
+    with open(file_path, 'r') as file:
+        lines = file.readlines()
+
+    with open(file_path, 'w') as file:
+        section_found = False
+        for line in lines:
+            if line.strip() == f'[{section}]':
+                section_found = True
+            if section_found and line.strip().startswith(key):
+                line = f'{key} = {new_value}\n'
+                section_found = False  # Reset the flag
+            file.write(line)
+
+
+def picroscope_authenticate_and_update_token(credentials_file):
+    """
+    Authentication and update service-account token for legacy picroscope environment. This updates the AWS credentials file
+    with the JWT token and updates it if it has <3 months before expiration. This function can be run as a cron job.
+    """
+    # Check if the JWT token exists and if it exists in the credentials file if it's expired.
+    # The credentials file section is [strapi] with `api_key` containing the jwt token, and `api_key_expires` containing
+    # the expiration date in ISO format.
+    config_file_path = os.path.expanduser(credentials_file)
+
+    config = configparser.ConfigParser()
+    with open(config_file_path, 'r') as f:
+        config.read_string(f.read())
+
+    assert 'strapi' in config, \
+        'Your AWS credentials file is missing a section [strapi], you may have the wrong version of the credentials file.'
+
+    token_exists = 'api_key' in config['strapi']
+    expire_exists = 'api_key_expires' in config['strapi']
+
+    if expire_exists:
+        expiration_str = config['strapi']['api_key_expires']
+        expiration_str = expiration_str.split(' ')[0] + ' ' + expiration_str.split(' ')[1]  # Remove timezone
+        expiration_date = datetime.datetime.fromisoformat(expiration_str)
+        days_remaining = (expiration_date - datetime.datetime.now()).days
+        print('Days remaining for token:', days_remaining)
+    else:
+        days_remaining = -1
+
+    # check if api_key_expires exists, if not, it's expired, else check if it has <90 days remaining on it
+    manual_refresh = not token_exists \
+                     or not expire_exists \
+                     or (datetime.datetime.fromisoformat(config['strapi']['api_key_expires']) - datetime.datetime.now()).days < 0
+    auto_refresh = (token_exists and expire_exists) \
+                   and (datetime.datetime.fromisoformat(config['strapi']['api_key_expires']) - datetime.datetime.now()).days < 90
+
+    if manual_refresh or auto_refresh:
+        token_data = authenticate_and_get_token() if manual_refresh else requests.get(url).json()
+        update_config_file(config_file_path, 'strapi', 'api_key', token_data['access_token'])
+        update_config_file(config_file_path, 'strapi', 'api_key_expires', token_data['expires_at'])
+        print(f'JWT token has been updated in {config_file_path}')
+    else:
+        print('JWT token is still valid, no action taken.')
+
+
+def parse_args():
+    """
+    Two commands are available:
+
+        # Authenticate and obtain a JWT service account token for braingeneerspy
+        python -m braingeneers.iot.messaging authenticate
+
+        # Authenticate and obtain a JWT service account token for picroscope specific environment
+        python -m braingeneers.iot.messaging authenticate picroscope
+    """
+    parser = argparse.ArgumentParser(description='JWT Service Account Token Management')
+    parser.add_argument('config', nargs='?', choices=['picroscope'], help='Picroscope specific JWT token configuration.')
+    parser.add_argument('--credentials', default='~/.aws/credentials', help='Path to the AWS credentials file, only used for picroscope authentication.')
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    if args.config == 'picroscope':
+        credentials_file = args.credentials
+        picroscope_authenticate_and_update_token(credentials_file)
+    else:
+        authenticate_and_get_token()
+
+
+if __name__ == '__main__':
+    main()

src/braingeneers/iot/messaging.py

@@ -12,6 +12,9 @@
 import json
 import braingeneers.iot.shadows as sh
 import pickle
+import importlib
+import argparse
+import datetime
 
 from typing import Callable, Tuple, List, Dict, Union
 from deprecated import deprecated
@@ -162,8 +165,9 @@ def __init__(self, name: str = None, credentials_file: (str, io.IOBase) = None,
         self._boto_iot_client = None
         self._boto_iot_data_client = None
         self._redis_client = None
+        self._jwt_service_account_token = None
 
-        self.shadow_interface = sh.DatabaseInteractor()
+        self.shadow_interface = sh.DatabaseInteractor(jwt_service_token=self.jwt_service_account_token)
 
         self._subscribed_data_streams = set()  # keep track of subscribed data streams
         self._subscribed_message_callback_map = {}  # keep track of subscribed message callbacks, key is regex, value is tuple of (callback, topic)
@@ -789,6 +793,38 @@ def redis_client(self) -> redis.Redis:
 
         return self._redis_client
 
+    @property
+    def jwt_service_account_token(self) -> str:
+        """ Lazy initialization of the JWT service account token. """
+        PACKAGE_NAME = "braingeneers.iot"
+        config_dir = os.path.join(importlib.resources.files(PACKAGE_NAME), 'service_account')
+        config_file = os.path.join(config_dir, 'config.json')
+
+        if self._jwt_service_account_token is None:
+            # Check if the JWT token exists
+            # This token is required for all operations that require web services.
+            # The token is a (json) dict of form {'access_token': '----', 'expires_at': '2024-11-07 23:39:42 UTC'}
+            os.makedirs(config_dir, exist_ok=True)
+
+            # Try to load an existing JWT token locally if it exists
+            if os.path.exists(config_file):
+                with open(config_file, 'r') as f:
+                    self._jwt_service_account_token = json.load(f)
+
+            if self._jwt_service_account_token is None:
+                raise PermissionError('JWT service account token not found, please generate one using: python -m braingeneers.iot.messaging authenticate')
+
+        # Check if the token is still valid, this happens on every access, but takes no action while it's still valid.
+        # If the token has less than 3 month left, refresh it, default tokens have 30 days at issuance.
+        expires_at = datetime.datetime.fromisoformat(self._jwt_service_account_token['expires_at'].replace(' UTC', ''))
+        if (expires_at - datetime.datetime.now()).days < 90:
+            GENERATE_TOKEN_URL = 'https://service-accounts.braingeneers.gi.ucsc.edu/generate_token'
+            self._jwt_service_account_token = requests.get(GENERATE_TOKEN_URL).json()
+            with open(config_file, 'w') as f:
+                json.dump(self._jwt_service_account_token, f)
+
+        return self._jwt_service_account_token
+
     def shutdown(self):
         """ Release resources and shutdown connections as needed. """
         if self.certs_temp_dir is not None: