Add basic JWT setup (#297)

d-j-hatton · web-flow · commit 160e08c024a3 · 2024-06-10T12:46:30.000+01:00
Add functionality to generate and use tokens for basic authorisation. Intend to extend
diff --git a/pyproject.toml b/pyproject.toml
@@ -60,8 +60,10 @@ server = [
     "jinja2",
     "numpy",
     "packaging",
+    "passlib",
     "pillow",
     "prometheus_client",
+    "python-jose[cryptography]",
     "readlif", # Specific to cryo-CLEM workflow
     "sqlmodel",
     "stomp-py<=8.1.0", # 8.1.1 (released 2024-04-06) doesn't work with our project
@@ -75,6 +77,7 @@ Documentation = "https://github.com/DiamondLightSource/python-murfey"
 GitHub = "https://github.com/DiamondLightSource/python-murfey"
 [project.scripts]
 murfey = "murfey.client:run"
+"murfey.add_user" = "murfey.cli.add_user:run"
 "murfey.create_db" = "murfey.cli.create_db:run"
 "murfey.db_sql" = "murfey.cli.murfey_db_sql:run"
 "murfey.decrypt_password" = "murfey.cli.decrypt_db_password:run"
diff --git a/src/murfey/cli/add_user.py b/src/murfey/cli/add_user.py
@@ -0,0 +1,40 @@
+import argparse
+import os
+
+from sqlmodel import Session, create_engine
+
+from murfey.server.auth import hash_password
+from murfey.server.config import get_machine_config
+from murfey.server.murfey_db import url
+from murfey.util.db import MurfeyUser as User
+
+
+def run():
+    parser = argparse.ArgumentParser(
+        description="Generate the necessary tables for the Murfey database"
+    )
+
+    parser.add_argument("-u", "--username", type=str, help="User name for new user")
+    parser.add_argument("-p", "--password", type=str, help="Password for new user")
+    parser.add_argument(
+        "-m",
+        "--microscope",
+        dest="microscope",
+        type=str,
+        default="",
+        help="Microscope as specified in the Murfey machine configuration",
+    )
+
+    args = parser.parse_args()
+    if args.microscope:
+        os.environ["BEAMLINE"] = args.microscope
+
+    new_user = User(
+        username=args.username, hashed_password=hash_password(args.password)
+    )
+    _url = url(get_machine_config())
+    engine = create_engine(_url)
+    with Session(engine) as murfey_db:
+        murfey_db.add(new_user)
+        murfey_db.commit()
+        murfey_db.close()
diff --git a/src/murfey/client/__init__.py b/src/murfey/client/__init__.py
@@ -12,6 +12,7 @@
 import time
 import webbrowser
 from datetime import datetime
+from functools import partial
 from pathlib import Path
 from queue import Queue
 from typing import List, Literal
@@ -41,6 +42,29 @@
 log = logging.getLogger("murfey.client")
 
 
+def read_config() -> configparser.ConfigParser:
+    config = configparser.ConfigParser()
+    try:
+        mcch = os.environ.get("MURFEY_CLIENT_CONFIG_HOME")
+        murfey_client_config_home = Path(mcch) if mcch else Path.home()
+        with open(murfey_client_config_home / ".murfey") as configfile:
+            config.read_file(configfile)
+    except FileNotFoundError:
+        log.warning(
+            f"Murfey client configuration file {murfey_client_config_home / '.murfey'} not found"
+        )
+    if "Murfey" not in config:
+        config["Murfey"] = {}
+    return config
+
+
+token = read_config()["Murfey"].get("token", "")
+
+requests.get = partial(requests.get, headers={"Authorization": f"Bearer {token}"})
+requests.post = partial(requests.post, headers={"Authorization": f"Bearer {token}"})
+requests.delete = partial(requests.delete, headers={"Authorization": f"Bearer {token}"})
+
+
 def _enable_webbrowser_in_cygwin():
     """Helper function to make webbrowser.open() work in CygWin"""
     if "cygwin" in platform.system().lower() and shutil.which("cygstart"):
@@ -317,20 +341,6 @@ def main_loop(
         time.sleep(15)
 
 
-def read_config() -> configparser.ConfigParser:
-    config = configparser.ConfigParser()
-    try:
-        mcch = os.environ.get("MURFEY_CLIENT_CONFIG_HOME")
-        murfey_client_config_home = Path(mcch) if mcch else Path.home()
-        with open(murfey_client_config_home / ".murfey") as configfile:
-            config.read_file(configfile)
-    except FileNotFoundError:
-        pass
-    if "Murfey" not in config:
-        config["Murfey"] = {}
-    return config
-
-
 def write_config(config: configparser.ConfigParser):
     mcch = os.environ.get("MURFEY_CLIENT_CONFIG_HOME")
     murfey_client_config_home = Path(mcch) if mcch else Path.home()
diff --git a/src/murfey/client/contexts/fib.py b/src/murfey/client/contexts/fib.py
@@ -10,9 +10,12 @@
 
 from murfey.client.context import Context
 from murfey.client.instance_environment import MurfeyInstanceEnvironment
+from murfey.util import authorised_requests
 
 logger = logging.getLogger("murfey.client.contexts.fib")
 
+requests.get, requests.post, requests.put, requests.delete = authorised_requests()
+
 
 class Lamella(NamedTuple):
     name: str
diff --git a/src/murfey/client/contexts/spa.py b/src/murfey/client/contexts/spa.py
@@ -15,10 +15,17 @@
     MurfeyID,
     MurfeyInstanceEnvironment,
 )
-from murfey.util import capture_get, capture_post, get_machine_config
+from murfey.util import (
+    authorised_requests,
+    capture_get,
+    capture_post,
+    get_machine_config,
+)
 
 logger = logging.getLogger("murfey.client.contexts.spa")
 
+requests.get, requests.post, requests.put, requests.delete = authorised_requests()
+
 
 class FoilHole(NamedTuple):
     session_id: int
diff --git a/src/murfey/client/contexts/spa_metadata.py b/src/murfey/client/contexts/spa_metadata.py
@@ -8,10 +8,12 @@
 from murfey.client.context import Context
 from murfey.client.contexts.spa import _get_grid_square_atlas_positions, _get_source
 from murfey.client.instance_environment import MurfeyInstanceEnvironment, SampleInfo
-from murfey.util import capture_post, get_machine_config
+from murfey.util import authorised_requests, capture_post, get_machine_config
 
 logger = logging.getLogger("murfey.client.contexts.spa_metadata")
 
+requests.get, requests.post, requests.put, requests.delete = authorised_requests()
+
 
 def _atlas_destination(
     environment: MurfeyInstanceEnvironment, source: Path, file_path: Path
diff --git a/src/murfey/client/contexts/tomo.py b/src/murfey/client/contexts/tomo.py
@@ -19,11 +19,13 @@
     MurfeyInstanceEnvironment,
     global_env_lock,
 )
-from murfey.util import capture_post, get_machine_config
+from murfey.util import authorised_requests, capture_post, get_machine_config
 from murfey.util.mdoc import get_block, get_global_data, get_num_blocks
 
 logger = logging.getLogger("murfey.client.contexts.tomo")
 
+requests.get, requests.post, requests.put, requests.delete = authorised_requests()
+
 
 class TiltInfoExtraction(NamedTuple):
     series: Callable[[Path], str]
diff --git a/src/murfey/client/tui/app.py b/src/murfey/client/tui/app.py
@@ -32,12 +32,23 @@
 from murfey.client.tui.status_bar import StatusBar
 from murfey.client.watchdir import DirWatcher
 from murfey.client.watchdir_multigrid import MultigridDirWatcher
-from murfey.util import capture_post, get_machine_config, set_default_acquisition_output
+from murfey.util import (
+    capture_post,
+    get_machine_config,
+    read_config,
+    set_default_acquisition_output,
+)
 
 log = logging.getLogger("murfey.tui.app")
 
 ReactiveType = TypeVar("ReactiveType")
 
+token = read_config()["Murfey"].get("token", "")
+
+requests.get = partial(requests.get, headers={"Authorization": f"Bearer {token}"})
+requests.post = partial(requests.post, headers={"Authorization": f"Bearer {token}"})
+requests.delete = partial(requests.delete, headers={"Authorization": f"Bearer {token}"})
+
 
 class MurfeyTUI(App):
     CSS_PATH = "controller.css"
diff --git a/src/murfey/client/tui/screens.py b/src/murfey/client/tui/screens.py
@@ -56,13 +56,19 @@
 )
 from murfey.client.rsync import RSyncer
 from murfey.client.tui.forms import FormDependency
-from murfey.util import capture_post, get_machine_config
+from murfey.util import capture_post, get_machine_config, read_config
 from murfey.util.models import PreprocessingParametersTomo, ProcessingParametersSPA
 
 log = logging.getLogger("murfey.tui.screens")
 
 ReactiveType = TypeVar("ReactiveType")
 
+token = read_config()["Murfey"].get("token", "")
+
+requests.get = partial(requests.get, headers={"Authorization": f"Bearer {token}"})
+requests.post = partial(requests.post, headers={"Authorization": f"Bearer {token}"})
+requests.delete = partial(requests.delete, headers={"Authorization": f"Bearer {token}"})
+
 
 def determine_default_destination(
     visit: str,
diff --git a/src/murfey/server/api.py b/src/murfey/server/api.py
@@ -9,7 +9,7 @@
 
 import packaging.version
 import sqlalchemy
-from fastapi import APIRouter, Request
+from fastapi import APIRouter, Depends, Request
 from fastapi.responses import FileResponse, HTMLResponse
 from ispyb.sqlalchemy import AutoProcProgram as ISPyBAutoProcProgram
 from ispyb.sqlalchemy import (
@@ -46,6 +46,7 @@
     sanitise,
     templates,
 )
+from murfey.server.auth import validate_token
 from murfey.server.config import from_file, settings
 from murfey.server.gain import Camera, prepare_eer_gain, prepare_gain
 from murfey.server.murfey_db import murfey_db
@@ -108,7 +109,7 @@
 
 machine_config = get_machine_config()
 
-router = APIRouter()
+router = APIRouter(dependencies=[Depends(validate_token)])
 
 
 # This will be the homepage for a given microscope.
diff --git a/src/murfey/server/auth/__init__.py b/src/murfey/server/auth/__init__.py
@@ -0,0 +1,67 @@
+import secrets
+from typing import Annotated
+
+from fastapi import Depends, HTTPException, status
+from fastapi.security import OAuth2PasswordBearer
+from jose import JWTError, jwt
+from passlib.context import CryptContext
+from sqlmodel import Session, create_engine, select
+
+from murfey.server.config import get_machine_config
+from murfey.server.murfey_db import url
+from murfey.util.db import MurfeyUser as User
+
+machine_config = get_machine_config()
+ALGORITHM = machine_config.auth_algorithm or "HS256"
+SECRET_KEY = machine_config.auth_key or secrets.token_hex(32)
+
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
+
+
+def verify_password(plain_password: str, hashed_password: str) -> bool:
+    return pwd_context.verify(plain_password, hashed_password)
+
+
+def hash_password(password: str) -> str:
+    return pwd_context.hash(password)
+
+
+try:
+    _url = url(get_machine_config())
+    engine = create_engine(_url)
+except Exception:
+    engine = None
+
+
+def validate_user(username: str, password: str) -> bool:
+    try:
+        with Session(engine) as murfey_db:
+            user = murfey_db.exec(select(User).where(User.username == username)).one()
+    except Exception:
+        return False
+    return verify_password(password, user.hashed_password)
+
+
+def check_user(username: str) -> bool:
+    try:
+        with Session(engine) as murfey_db:
+            users = murfey_db.exec(select(User)).all()
+    except Exception:
+        return False
+    return username in [u.username for u in users]
+
+
+async def validate_token(token: Annotated[str, Depends(oauth2_scheme)]):
+    try:
+        decoded_data = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
+        if not check_user(decoded_data.get("user")):
+            raise JWTError
+    except JWTError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Could not validate credentials",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    return None
diff --git a/src/murfey/server/auth/api.py b/src/murfey/server/auth/api.py
@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from logging import getLogger
+
+from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi.security import OAuth2PasswordRequestForm
+from jose import jwt
+from pydantic import BaseModel
+from typing_extensions import Annotated
+
+from murfey.server.auth import ALGORITHM, SECRET_KEY, validate_user
+
+logger = getLogger("murfey.server.auth.api")
+
+router = APIRouter()
+
+
+class Token(BaseModel):
+    access_token: str
+    token_type: str
+
+
+def create_access_token(data: dict):
+    to_encode = data.copy()
+
+    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
+    return encoded_jwt
+
+
+@router.post("/token")
+def generate_token(
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+) -> Token:
+    validated = validate_user(form_data.username, form_data.password)
+    if not validated:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+    access_token = create_access_token(
+        data={"user": form_data.username},
+    )
+    return Token(access_token=access_token, token_type="bearer")
diff --git a/src/murfey/server/config.py b/src/murfey/server/config.py
@@ -47,6 +47,8 @@ class MachineConfig(BaseModel):
     upstream_data_download_directory: Optional[Path] = None  # Set by microscope config
     upstream_data_tiff_locations: List[str] = ["processed"]  # Location of CLEM TIFFs
     failure_queue: str = ""
+    auth_key: str = ""
+    auth_algorithm: str = ""
 
 
 def from_file(config_file_path: Path, microscope: str) -> MachineConfig:
diff --git a/src/murfey/server/demo_api.py b/src/murfey/server/demo_api.py
diff --git a/src/murfey/server/main.py b/src/murfey/server/main.py
diff --git a/src/murfey/util/__init__.py b/src/murfey/util/__init__.py
diff --git a/src/murfey/util/db.py b/src/murfey/util/db.py
diff --git a/tests/server/api/test_movies.py b/tests/server/api/test_movies.py
diff --git a/tests/server/test_main.py b/tests/server/test_main.py
