Merge pull request #198 from Clinical-Genomics/replace-custom-oauth-header (patch)

karlnyr · web-flow · commit 1ecf105c6f49 · 2025-03-12T10:30:36.000+01:00
### Changed

- Replace custom oauth header
diff --git a/microSALT/utils/pubmlst/authentication.py b/microSALT/utils/pubmlst/authentication.py
@@ -1,22 +1,33 @@
 import json
 import os
 from datetime import datetime, timedelta
+
 from dateutil import parser
 from rauth import OAuth1Session
+
 from microSALT import logger
-from microSALT.utils.pubmlst.helpers import BASE_API, save_session_token, load_auth_credentials, get_path, folders_config, credentials_path_key, pubmlst_session_credentials_file_name 
 from microSALT.utils.pubmlst.exceptions import (
     PUBMLSTError,
     SessionTokenRequestError,
     SessionTokenResponseError,
 )
+from microSALT.utils.pubmlst.helpers import (
+    BASE_API,
+    credentials_path_key,
+    folders_config,
+    get_path,
+    load_auth_credentials,
+    pubmlst_session_credentials_file_name,
+    save_session_token,
+)
 
 session_token_validity = 12  # 12-hour validity
 session_expiration_buffer = 60  # 60-second buffer
 
+
 def get_new_session_token(db: str):
     """Request a new session token using all credentials for a specific database."""
-    logger.debug("Fetching a new session token for database '{db}'...")
+    logger.debug(f"Fetching a new session token for database '{db}'...")
 
     try:
         consumer_key, consumer_secret, access_token, access_secret = load_auth_credentials()
@@ -30,8 +41,8 @@ def get_new_session_token(db: str):
             access_token_secret=access_secret,
         )
 
-        response = session.get(url, headers={"User-Agent": "BIGSdb downloader"})
-        logger.debug("Response Status Code: {status_code}")
+        response = session.get(url, headers={"User-Agent": "BIGSdb API downloader"})
+        logger.debug(f"Response Status Code: {response.status_code}")
 
         if response.ok:
             try:
@@ -52,23 +63,23 @@ def get_new_session_token(db: str):
             except (ValueError, KeyError) as e:
                 raise SessionTokenResponseError(db, f"Invalid response format: {str(e)}")
         else:
-            raise SessionTokenRequestError(
-                db, response.status_code, response.text
-            )
+            raise SessionTokenRequestError(db, response.status_code, response.text)
 
     except PUBMLSTError as e:
         logger.error(f"Error during token fetching: {e}")
         raise
     except Exception as e:
         logger.error(f"Unexpected error: {e}")
-        raise PUBMLSTError(f"Unexpected error while fetching session token for database '{db}': {e}")
+        raise PUBMLSTError(
+            f"Unexpected error while fetching session token for database '{db}': {e}"
+        )
+
 
 def load_session_credentials(db: str):
     """Load session token from file for a specific database."""
     try:
         credentials_file = os.path.join(
-            get_path(folders_config, credentials_path_key),
-            pubmlst_session_credentials_file_name
+            get_path(folders_config, credentials_path_key), pubmlst_session_credentials_file_name
         )
 
         if not os.path.exists(credentials_file):
@@ -83,7 +94,9 @@ def load_session_credentials(db: str):
 
         db_session_data = all_sessions.get("databases", {}).get(db)
         if not db_session_data:
-            logger.debug(f"No session token found for database '{db}'. Fetching a new session token.")
+            logger.debug(
+                f"No session token found for database '{db}'. Fetching a new session token."
+            )
             return get_new_session_token(db)
 
         expiration = parser.parse(db_session_data.get("expiration", ""))
@@ -94,7 +107,9 @@ def load_session_credentials(db: str):
 
             return session_token, session_secret
 
-        logger.debug(f"Session token for database '{db}' has expired. Fetching a new session token.")
+        logger.debug(
+            f"Session token for database '{db}' has expired. Fetching a new session token."
+        )
         return get_new_session_token(db)
 
     except PUBMLSTError as e:
@@ -103,4 +118,3 @@ def load_session_credentials(db: str):
     except Exception as e:
         logger.error(f"Unexpected error: {e}")
         raise PUBMLSTError(f"Unexpected error while loading session token for database '{db}': {e}")
-
diff --git a/microSALT/utils/pubmlst/client.py b/microSALT/utils/pubmlst/client.py
@@ -1,68 +1,80 @@
-import requests
 from urllib.parse import urlencode
+
+import requests
+from rauth import OAuth1Session
+
+from microSALT import logger
+from microSALT.utils.pubmlst.authentication import load_session_credentials
+from microSALT.utils.pubmlst.constants import HTTPMethod, RequestType, ResponseHandler
+from microSALT.utils.pubmlst.exceptions import PUBMLSTError, SessionTokenRequestError
 from microSALT.utils.pubmlst.helpers import (
     BASE_API,
-    generate_oauth_header,
     load_auth_credentials,
-    parse_pubmlst_url
+    parse_pubmlst_url,
 )
-from microSALT.utils.pubmlst.constants import RequestType, HTTPMethod, ResponseHandler
-from microSALT.utils.pubmlst.exceptions import PUBMLSTError, SessionTokenRequestError
-from microSALT.utils.pubmlst.authentication import load_session_credentials
-from microSALT import logger
+
 
 class PubMLSTClient:
     """Client for interacting with the PubMLST authenticated API."""
 
     def __init__(self):
         """Initialize the PubMLST client."""
         try:
-            self.consumer_key, self.consumer_secret, self.access_token, self.access_secret = load_auth_credentials()
+            self.consumer_key, self.consumer_secret, self.access_token, self.access_secret = (
+                load_auth_credentials()
+            )
             self.database = "pubmlst_test_seqdef"
             self.session_token, self.session_secret = load_session_credentials(self.database)
         except PUBMLSTError as e:
             logger.error(f"Failed to initialize PubMLST client: {e}")
             raise
 
-
     @staticmethod
     def parse_pubmlst_url(url: str):
         """
         Wrapper for the parse_pubmlst_url function.
         """
         return parse_pubmlst_url(url)
 
-
-    def _make_request(self, request_type: RequestType, method: HTTPMethod, url: str, db: str = None, response_handler: ResponseHandler = ResponseHandler.JSON):
-        """ Handle API requests."""     
+    def _make_request(
+        self,
+        request_type: RequestType,
+        method: HTTPMethod,
+        url: str,
+        db: str = None,
+        response_handler: ResponseHandler = ResponseHandler.JSON,
+    ):
+        """Handle API requests."""
         try:
-            if db:
-                session_token, session_secret = load_session_credentials(db)
-            else:
-                session_token, session_secret = self.session_token, self.session_secret
-            
             if request_type == RequestType.AUTH:
-                headers = {
-                    "Authorization": generate_oauth_header(url, self.consumer_key, self.consumer_secret, self.access_token, self.access_secret)
-                }
+                access_token = self.access_token
+                access_secret = self.access_secret
+                log_database = "authentication"
             elif request_type == RequestType.DB:
-                headers = {
-                    "Authorization": generate_oauth_header(url, self.consumer_key, self.consumer_secret, session_token, session_secret)
-                }
+                access_token, access_secret = load_session_credentials(db or self.database)
+                log_database = db or self.database
             else:
                 raise ValueError(f"Unsupported request type: {request_type}")
 
-            if method == HTTPMethod.GET:
-                response = requests.get(url, headers=headers)
-            elif method == HTTPMethod.POST:
-                response = requests.post(url, headers=headers)
-            elif method == HTTPMethod.PUT:
-                response = requests.put(url, headers=headers)
-            else:
-                raise ValueError(f"Unsupported HTTP method: {method}")
-            
+            logger.info(f"Making request to {url} for database {log_database}")
+            logger.info(f"Using session token: {access_token[:6]}****")
+
+            session = OAuth1Session(
+                consumer_key=self.consumer_key,
+                consumer_secret=self.consumer_secret,
+                access_token=access_token,
+                access_token_secret=access_secret,
+            )
+
+            headers = {"User-Agent": "BIGSdb API Client"}
+            response = session.request(method.value, url, headers=headers)
+
+            if response.status_code == 401:
+                logger.error(f"401 Unauthorized: {response.text}")
+
             response.raise_for_status()
-            
+
+            # Process the response based on the requested handler type
             if response_handler == ResponseHandler.CONTENT:
                 return response.content
             elif response_handler == ResponseHandler.TEXT:
@@ -73,44 +85,51 @@ def _make_request(self, request_type: RequestType, method: HTTPMethod, url: str,
                 raise ValueError(f"Unsupported response handler: {response_handler}")
 
         except requests.exceptions.HTTPError as e:
-            raise SessionTokenRequestError(db or self.database, e.response.status_code, e.response.text) from e
+            raise SessionTokenRequestError(
+                db or self.database, e.response.status_code, e.response.text
+            ) from e
         except requests.exceptions.RequestException as e:
             logger.error(f"Request failed: {e}")
             raise PUBMLSTError(f"Request failed: {e}") from e
         except Exception as e:
             logger.error(f"Unexpected error during request: {e}")
             raise PUBMLSTError(f"An unexpected error occurred: {e}") from e
 
-
     def query_databases(self):
         """Query available PubMLST databases."""
         url = f"{BASE_API}/db"
-        return self._make_request(RequestType.DB, HTTPMethod.GET, url, response_handler=ResponseHandler.JSON)
-
+        return self._make_request(
+            RequestType.DB, HTTPMethod.GET, url, response_handler=ResponseHandler.JSON
+        )
 
     def download_locus(self, db: str, locus: str, **kwargs):
         """Download locus sequence files."""
         base_url = f"{BASE_API}/db/{db}/loci/{locus}/alleles_fasta"
         query_string = urlencode(kwargs)
         url = f"{base_url}?{query_string}" if query_string else base_url
-        return self._make_request(RequestType.DB, HTTPMethod.GET, url, db=db, response_handler=ResponseHandler.TEXT)
-
+        return self._make_request(
+            RequestType.DB, HTTPMethod.GET, url, db=db, response_handler=ResponseHandler.TEXT
+        )
 
     def download_profiles_csv(self, db: str, scheme_id: int):
         """Download MLST profiles in CSV format."""
         if not scheme_id:
             raise ValueError("Scheme ID is required to download profiles CSV.")
         url = f"{BASE_API}/db/{db}/schemes/{scheme_id}/profiles_csv"
-        return self._make_request(RequestType.DB, HTTPMethod.GET, url, db=db, response_handler=ResponseHandler.TEXT)
-
+        return self._make_request(
+            RequestType.DB, HTTPMethod.GET, url, db=db, response_handler=ResponseHandler.TEXT
+        )
 
     def retrieve_scheme_info(self, db: str, scheme_id: int):
         """Retrieve information about a specific MLST scheme."""
         url = f"{BASE_API}/db/{db}/schemes/{scheme_id}"
-        return self._make_request(RequestType.DB, HTTPMethod.GET, url, db=db, response_handler=ResponseHandler.JSON)
-
+        return self._make_request(
+            RequestType.DB, HTTPMethod.GET, url, db=db, response_handler=ResponseHandler.JSON
+        )
 
     def list_schemes(self, db: str):
         """List available MLST schemes for a specific database."""
         url = f"{BASE_API}/db/{db}/schemes"
-        return self._make_request(RequestType.DB, HTTPMethod.GET, url, db=db, response_handler=ResponseHandler.JSON)
+        return self._make_request(
+            RequestType.DB, HTTPMethod.GET, url, db=db, response_handler=ResponseHandler.JSON
+        )
diff --git a/microSALT/utils/pubmlst/helpers.py b/microSALT/utils/pubmlst/helpers.py
@@ -1,15 +1,24 @@
-import os
 import base64
 import hashlib
-import json
 import hmac
+import json
+import os
 import time
 from pathlib import Path
 from urllib.parse import quote_plus, urlencode
+
 from werkzeug.exceptions import NotFound
+
 from microSALT import app, logger
-from microSALT.utils.pubmlst.exceptions import PUBMLSTError, PathResolutionError, CredentialsFileNotFound, InvalidCredentials, SaveSessionError, InvalidURLError
 from microSALT.utils.pubmlst.constants import Encoding, url_map
+from microSALT.utils.pubmlst.exceptions import (
+    CredentialsFileNotFound,
+    InvalidCredentials,
+    InvalidURLError,
+    PathResolutionError,
+    PUBMLSTError,
+    SaveSessionError,
+)
 
 BASE_WEB = "https://pubmlst.org/bigsdb"
 BASE_API = "https://rest.pubmlst.org"
@@ -21,6 +30,7 @@
 pubmlst_config = app.config["pubmlst"]
 folders_config = app.config["folders"]
 
+
 def get_path(config, config_key: str):
     """Get and expand the file path from the configuration."""
     try:
@@ -41,8 +51,7 @@ def load_auth_credentials():
     """Load client ID, client secret, access token, and access secret from credentials file."""
     try:
         credentials_file = os.path.join(
-            get_path(folders_config, credentials_path_key),
-            pubmlst_auth_credentials_file_name
+            get_path(folders_config, credentials_path_key), pubmlst_auth_credentials_file_name
         )
 
         if not os.path.exists(credentials_file):
@@ -81,37 +90,7 @@ def load_auth_credentials():
         raise
     except Exception as e:
         raise PUBMLSTError("An unexpected error occurred while loading credentials: {e}")
-    
-
-def generate_oauth_header(url: str, oauth_consumer_key: str, oauth_consumer_secret: str, oauth_token: str, oauth_token_secret: str):
-    """Generate the OAuth1 Authorization header."""
-    oauth_timestamp = str(int(time.time()))
-    oauth_nonce = base64.urlsafe_b64encode(os.urandom(32)).decode(Encoding.UTF8.value).strip("=")
-    oauth_signature_method = "HMAC-SHA1"
-    oauth_version = "1.0"
-
-    oauth_params = {
-        "oauth_consumer_key": oauth_consumer_key,
-        "oauth_token": oauth_token,
-        "oauth_signature_method": oauth_signature_method,
-        "oauth_timestamp": oauth_timestamp,
-        "oauth_nonce": oauth_nonce,
-        "oauth_version": oauth_version,
-    }    
-
-    params_encoded = urlencode(sorted(oauth_params.items()))
-    base_string = f"GET&{quote_plus(url)}&{quote_plus(params_encoded)}"
-    signing_key = f"{oauth_consumer_secret}&{oauth_token_secret}"
-
-    hashed = hmac.new(signing_key.encode(Encoding.UTF8.value), base_string.encode(Encoding.UTF8.value), hashlib.sha1)
-    oauth_signature = base64.b64encode(hashed.digest()).decode(Encoding.UTF8.value)
-
-    oauth_params["oauth_signature"] = oauth_signature
-
-    auth_header = "OAuth " + ", ".join(
-        [f'{quote_plus(k)}="{quote_plus(v)}"' for k, v in oauth_params.items()]
-    )
-    return auth_header
+
 
 def save_session_token(db: str, token: str, secret: str, expiration_date: str):
     """Save session token, secret, and expiration to a JSON file for the specified database."""
@@ -123,8 +102,7 @@ def save_session_token(db: str, token: str, secret: str, expiration_date: str):
         }
 
         credentials_file = os.path.join(
-            get_path(folders_config, credentials_path_key),
-            pubmlst_session_credentials_file_name
+            get_path(folders_config, credentials_path_key), pubmlst_session_credentials_file_name
         )
 
         if os.path.exists(credentials_file):
@@ -141,16 +119,15 @@ def save_session_token(db: str, token: str, secret: str, expiration_date: str):
         with open(credentials_file, "w") as f:
             json.dump(all_sessions, f, indent=4)
 
-        logger.debug(
-            f"Session token for database '{db}' saved to '{credentials_file}'."
-        )
+        logger.debug(f"Session token for database '{db}' saved to '{credentials_file}'.")
     except (IOError, OSError) as e:
         raise SaveSessionError(db, f"I/O error: {e}")
     except ValueError as e:
         raise SaveSessionError(db, f"Invalid data format: {e}")
     except Exception as e:
         raise SaveSessionError(db, f"Unexpected error: {e}")
 
+
 def parse_pubmlst_url(url: str):
     """
     Match a URL against the URL map and return extracted parameters.
