diff --git a/Makefile b/Makefile index cdcfd6e..e1b01be 100644 --- a/Makefile +++ b/Makefile @@ -15,5 +15,7 @@ run: build --name $(SERVER_CONTAINER_NAME) \ -e DB_NAME=$(DB_NAME) \ -e MONGO_URL=$(MONGO_URL) \ - --network=host \ + -e JWT_SECRET_KEY=$(JWT_SECRET_KEY) \ + -p 5000:5000 \ + -p 8765:8765 \ $(SERVER_IMAGE) diff --git a/app/config/jwt_config.py b/app/config/jwt_config.py new file mode 100644 index 0000000..f758c28 --- /dev/null +++ b/app/config/jwt_config.py @@ -0,0 +1,15 @@ +import os +from quart_jwt_extended import JWTManager + +# Initialize JWTManager outside of any function to ensure it's a singleton +# It will be initialized with the app object later in server.py +jwt = JWTManager() + +def configure_jwt(app): + """Configures JWT settings for the Quart app.""" + app.config["JWT_SECRET_KEY"] = os.getenv("JWT_SECRET_KEY", "your-super-secret-key-that-should-be-in-env") + app.config["JWT_ACCESS_TOKEN_EXPIRES"] = 3600 # 1 hour + # You might need to set a custom error handler for unauthorized access + # @app.errorhandler(401) + # def unauthorized(error): + # return jsonify({"msg": "Missing or invalid token"}), 401 \ No newline at end of file diff --git a/app/internal/auth_wrappers.py b/app/internal/auth_wrappers.py new file mode 100644 index 0000000..db36d20 --- /dev/null +++ b/app/internal/auth_wrappers.py @@ -0,0 +1,70 @@ +# internal/auth_wrappers.py +import os +import asyncio +from uuid import uuid4 +from typing import Optional, List, Dict, Any +from internal.db_handler import MongoHandler +from internal.types import User, UserRoles + +DB_NAME = os.getenv("DB_NAME", "default_db") +MONGO_URL = os.getenv("MONGO_URL", "mongodb://10.10.202.4:27017/") + +USER_DB_COLLECTION_NAME = "users" + +class UserDbController: + def __init__(self): + self.db_h = MongoHandler(DB_NAME, USER_DB_COLLECTION_NAME, MONGO_URL) + + async def create_user(self, user_data: Dict[str, Any]) -> Optional[User]: + try: + if not user_data.get("_id"): + user_data['_id'] = str(uuid4()) + + inserted_id = None + async with self.db_h as db: + insert_result = await db.insert_one(user_data) + inserted_id = insert_result.inserted_id + + if inserted_id: + query = {"_id": inserted_id} + inserted_doc = None + async with self.db_h as db: + inserted_doc = await db.find_one(query) + if inserted_doc: + return User.from_dict(inserted_doc) + return None + except Exception as e: + print(f"Create user failed: {e}") + return None + + async def find_user(self, query: Dict[str, Any]) -> Optional[User]: + try: + found_doc = None + async with self.db_h as db: + found_doc = await db.find_one(query) + if found_doc: + return User.from_dict(found_doc) + return None + except Exception as e: + print(f"Find user failed: {e}") + return None + + async def update_user(self, query: Dict[str, Any], update_data: Dict[str, Any]) -> Optional[int]: + try: + update_result = None + async with self.db_h as db: + update_result = await db.update_one(query, update_data) + return update_result.modified_count + except Exception as e: + print(f"Update user failed: {e}") + return None + + async def delete_user(self, query: Dict[str, Any]) -> Optional[int]: + try: + delete_result = None + async with self.db_h as db: + delete_result = await db.delete_one(query) + return delete_result.deleted_count + except Exception as e: + print(f"Delete user failed: {e}") + return None \ No newline at end of file diff --git a/app/internal/types.py b/app/internal/types.py index e187335..16d76e9 100644 --- a/app/internal/types.py +++ b/app/internal/types.py @@ -1,3 +1,4 @@ +# internal/types.py from typing import Optional, List, Dict, Any from enum import Enum @@ -20,7 +21,7 @@ class TalkgroupTag: # Add a method to convert to a dictionary, useful for sending as JSON def to_dict(self) -> Dict[str, Any]: - return {"talkgroup": self.talkgroup, "tagDec": self.tagDec} + return {"talkgroup": self.talkgroup, "tagDec": self.talkgroup} class DiscordId: @@ -189,4 +190,48 @@ class ActiveClient: def __init__(self, websocket= None, active_token:DiscordId=None): self.active_token = active_token - self.websocket = websocket \ No newline at end of file + self.websocket = websocket + +class UserRoles(str, Enum): + ADMIN = "admin" + MOD = "mod" + USER = "user" + +class User: + """ + A data model for a User entry. + """ + def __init__(self, + _id: str, + username: str, + password_hash: str, + role: UserRoles, + api_key: Optional[str] = None): + self._id: str = str(_id) + self.username: str = username + self.password_hash: str = password_hash + self.role: UserRoles = role + self.api_key: Optional[str] = api_key + + def __repr__(self) -> str: + return (f"User(_id='{self._id}', username='{self.username}', role='{self.role.value}', " + f"api_key={'' if self.api_key else 'None'})") + + def to_dict(self) -> Dict[str, Any]: + return { + "_id": self._id, + "username": self.username, + "password_hash": self.password_hash, + "role": self.role.value, + "api_key": self.api_key, + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "User": + return cls( + _id=data.get("_id"), + username=data.get("username", ""), + password_hash=data.get("password_hash", ""), + role=UserRoles(data.get("role", UserRoles.USER.value)), + api_key=data.get("api_key", None), + ) \ No newline at end of file diff --git a/app/routers/auth.py b/app/routers/auth.py new file mode 100644 index 0000000..73e7aee --- /dev/null +++ b/app/routers/auth.py @@ -0,0 +1,130 @@ +import functools +from quart import Blueprint, jsonify, request, current_app, abort +from werkzeug.security import generate_password_hash, check_password_hash +from quart_jwt_extended import create_access_token, jwt_required, get_jwt_identity +from internal.auth_wrappers import UserDbController +from internal.types import UserRoles +from uuid import uuid4 + +# Import the centralized JWTManager instance +from config.jwt_config import jwt as jwt_manager_instance # Renamed to avoid confusion with jwt_required + + +auth_bp = Blueprint('auth', __name__) + +# Decorator for role-based access control +def role_required(required_role: UserRoles): + def wrapper(fn): + @functools.wraps(fn) + @jwt_required + async def decorated_view(*args, **kwargs): + current_user_identity = get_jwt_identity() + user_id = current_user_identity['id'] + + # Make a DB call to get the user and their role + user = await current_app.user_db_h.find_user({"_id": user_id}) + + if not user: + abort(401, "User not found or invalid token.") # User corresponding to token not found + + user_role = user.role # Get the role from the fetched user object + + role_order = {UserRoles.USER: 0, UserRoles.MOD: 1, UserRoles.ADMIN: 2} + + if role_order[user_role] < role_order[required_role]: + abort(403, "Permission denied: Insufficient role.") + + # REMOVE current_app.ensure_sync() here + return await fn(*args, **kwargs) # Directly await the original async function + + return decorated_view + return wrapper + +@auth_bp.route('/register', methods=['POST']) +async def register_user(): + data = await request.get_json() + username = data.get('username') + password = data.get('password') + role = data.get('role', UserRoles.USER.value) # Default to 'user' role + + if not username or not password: + abort(400, "Username and password are required") + + existing_user = await current_app.user_db_h.find_user({"username": username}) + if existing_user: + abort(409, "Username already exists") + + hashed_password = generate_password_hash(password) + + try: + user_role = UserRoles(role) + except ValueError: + abort(400, f"Invalid role: {role}. Must be one of {list(UserRoles)}") + + user_data = { + "username": username, + "password_hash": hashed_password, + "role": user_role.value + } + + new_user = await current_app.user_db_h.create_user(user_data) + + if new_user: + return jsonify({"message": "User registered successfully", "user_id": new_user._id}), 201 + else: + abort(500, "Failed to register user") + +@auth_bp.route('/login', methods=['POST']) +async def login_user(): + data = await request.get_json() + username = data.get('username') + password = data.get('password') + + if not username or not password: + abort(400, "Username and password are required") + + user = await current_app.user_db_h.find_user({"username": username}) + + if not user or not check_password_hash(user.password_hash, password): + abort(401, "Invalid credentials") + + access_token = create_access_token(identity={"id": user._id, "username": user.username}) + return jsonify(access_token=access_token), 200 + +@auth_bp.route('/generate_api_key', methods=['POST']) +@jwt_required +async def generate_api_key(): + current_user_identity = get_jwt_identity() + user_id = current_user_identity['id'] + user_role = current_user_identity['role'] + + if user_role not in [UserRoles.ADMIN.value, UserRoles.MOD.value]: + target_user_id = (await request.get_json()).get('user_id', user_id) + if target_user_id != user_id: + abort(403, "Permission denied: You can only generate an API key for your own account.") + else: + target_user_id = (await request.get_json()).get('user_id', user_id) + + + new_api_key = str(uuid4()) + updated_count = await current_app.user_db_h.update_user( + {"_id": target_user_id}, + {"$set": {"api_key": new_api_key}} + ) + + if updated_count: + return jsonify({"message": f"API key generated for user {target_user_id}", "api_key": new_api_key}), 200 + else: + abort(404, f"User {target_user_id} not found or unable to update API key.") + +@auth_bp.route('/admin_only', methods=['GET']) +@jwt_required +@role_required(UserRoles.ADMIN) +async def admin_only_route(): + return jsonify({"message": "Welcome, Admin!"}), 200 + +@auth_bp.route('/mod_or_admin_only', methods=['GET']) +@jwt_required +@role_required(UserRoles.MOD) +async def mod_or_admin_only_route(): + return jsonify({"message": "Welcome, Mod or Admin!"}), 200 \ No newline at end of file diff --git a/app/server.py b/app/server.py index beea133..4f88025 100644 --- a/app/server.py +++ b/app/server.py @@ -1,3 +1,4 @@ +# server.py import asyncio import websockets import json @@ -6,7 +7,12 @@ from quart import Quart, jsonify, request from routers.systems import systems_bp from routers.nodes import nodes_bp, register_client, unregister_client from routers.bot import bot_bp +from routers.auth import auth_bp # ONLY import auth_bp, not jwt instance from auth.py from internal.db_wrappers import SystemDbController, DiscordIdDbController +from internal.auth_wrappers import UserDbController +# Import the JWTManager instance and its configuration function +from config.jwt_config import jwt, configure_jwt # Import the actual jwt instance and the config function + # --- WebSocket Server Components --- # Dictionary to store active clients: {client_id: websocket} @@ -57,6 +63,12 @@ app.active_clients = active_clients # Create and attach the DB wrappers app.sys_db_h = SystemDbController() app.d_id_db_h = DiscordIdDbController() +app.user_db_h = UserDbController() + +# Configure JWT settings and initialize the JWTManager instance with the app +configure_jwt(app) +jwt.init_app(app) # Crucial: This initializes the global 'jwt' instance with your app + @app.before_serving async def startup_websocket_server(): @@ -86,6 +98,7 @@ async def shutdown_websocket_server(): app.register_blueprint(systems_bp, url_prefix="/systems") app.register_blueprint(nodes_bp, url_prefix="/nodes") app.register_blueprint(bot_bp, url_prefix="/bots") +app.register_blueprint(auth_bp, url_prefix="/auth") # Register the auth blueprint @app.route('/') async def index(): @@ -93,14 +106,10 @@ async def index(): # --- Main Execution --- if __name__ == "__main__": - # Quart's app.run() will start the asyncio event loop and manage it. - # The @app.before_serving decorator ensures the websocket server starts within that loop. - # We removed asyncio.run(main()) and the main() function itself. print("Starting Quart API server...") app.run( host="0.0.0.0", port=5000, debug=False # Set to True for development ) - print("Quart API server stopped.") - + print("Quart API server stopped.") \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index f173fc1..52bcfea 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,6 @@ websockets quart motor -fastapi \ No newline at end of file +fastapi +quart-jwt-extended +passlib[bcrypt] \ No newline at end of file