Added module server

Author: TBWanderer

src/server/__init__.py

@@ -0,0 +1 @@
+from .websock import init

src/server/commands.py

@@ -0,0 +1,105 @@
+import asyncio
+import sys
+import uuid
+
+from .users import get_users, save_users, save_user
+from .crypto import generate_key_pair
+
+
+async def start_command_reader():
+    while True:
+        # This creates a future that will be set when stdin has data
+        line = await asyncio.to_thread(sys.stdin.readline)
+        line = line.strip()
+
+        if line:
+            await process_command(line)
+
+
+async def process_command(command: str) -> None:
+    cmd_parts = command.split()
+    if not cmd_parts:
+        return
+
+    cmd = cmd_parts[0].lower()
+
+    if cmd == "help":
+        print("Available commands:")
+        print("  new - Create new user")
+        print("  list - List all registered users")
+        print("  del <uuid> - Delete a user by UUID")
+        print("  exit - Exit the server")
+        print("  help - Show this help message")
+
+    elif cmd == "new":
+        if len(cmd_parts) != 1:
+            print("Usage: new")
+            return
+
+        await generate_user_keys()
+
+    elif cmd == "list":
+        users = get_users()
+        if not users:
+            print("No users registered.")
+            return
+
+        print("Registered users:")
+        for user_id, user_data in users.items():
+            print(f"  UUID: {user_id}")
+
+    elif cmd == "del":
+        if len(cmd_parts) < 2:
+            print("Usage: del <uuid>")
+            return
+
+        user_id = cmd_parts[1]
+        users = get_users()
+        if user_id in users:
+            del users[user_id]
+            save_users(users)
+            print(f"User {user_id} deleted.")
+        else:
+            print(f"User {user_id} not found.")
+
+    elif cmd == "exit":
+        print("Shutting down server...")
+        # Signal to stop the server
+        asyncio.get_event_loop().stop()
+        sys.exit()
+
+    else:
+        print(f"Unknown command: {cmd}")
+        print("Type 'help' for available commands.")
+
+
+async def generate_user_keys() -> None:
+    import get_from_env
+
+    # Generate server keys
+    server_secret, server_public = generate_key_pair()
+
+    # Generate client keys
+    client_secret, client_public = generate_key_pair()
+
+    # Generate UUID for the user
+    user_id = str(uuid.uuid4())
+    server_conf = {
+        "secret": server_secret,
+        "public_key": client_public
+    }
+    user_conf = {
+        "url": f"ws://{get_from_env.get_addr()}:{get_from_env.get_port()}",
+        "uuid": user_id,
+        "secret": client_secret,
+        "public": server_public
+    }
+
+    # Save to users.json
+    users = get_users()
+    users[user_id] = server_conf
+    save_users(users)
+
+    save_user(user_id, user_conf)
+
+    print(f"New user UUID: {user_id}")

src/server/crypto.py

@@ -0,0 +1,22 @@
+import base64
+import ecies
+
+
+def generate_key_pair():
+    key = ecies.utils.generate_key()
+    secret = key.secret
+    public = key.public_key.format(True)
+
+    # Convert to bytes and then base64 for storage
+    secret_b64 = base64.b64encode(secret).decode("utf-8")
+    public_b64 = base64.b64encode(public).decode("utf-8")
+
+    return secret_b64, public_b64
+
+
+def decrypt_data(private_key, encrypted_data):
+    return ecies.decrypt(private_key, encrypted_data).decode("utf-8")
+
+
+def encrypt_data(public_key, data):
+    return ecies.encrypt(public_key, data)

src/server/get_from_env.py

@@ -0,0 +1,20 @@
+import os
+import logging
+
+def get_addr() -> str:
+    addr = os.getenv("SWUC_SERVER_ADDR")
+    if not addr:
+        return "localhost"
+    else:
+        return addr
+
+
+def get_port() -> int:
+    port = os.getenv("SWUC_SERVER_PORT")
+    try:
+        port = int(port) if port is not None else 8765
+    except ValueError:
+        logging.warning(f"Incorrect port value: '{port}', using default 8765")
+        port = 8765
+
+    return port

src/server/users.py

@@ -0,0 +1,33 @@
+import json
+import os
+from logging import warning
+
+def get_users() -> dict:
+    try:
+        if not os.path.exists("users.json"):
+            with open("users.json", "w") as f:
+                json.dump({}, f)
+
+        with open("users.json", "r") as users_file:
+            return json.load(users_file)
+    except json.JSONDecodeError:
+        return {}
+    except Exception as e:
+        warning(f"Error loading users: {repr(e)}")
+        return {}
+
+def save_users(users: dict) -> None:
+    try:
+        with open("users.json", "w") as users_file:
+            json.dump(users, users_file, indent=2)
+    except Exception as e:
+        warning(f"Error saving users: {repr(e)}")
+
+def save_user(uuid: str, user: dict) -> None:
+    if not os.path.exists("users"):
+        os.makedirs("users")
+    try:
+        with open(f"users/{uuid}.json", "w") as users_file:
+            json.dump(user, users_file, indent=2)
+    except Exception as e:
+        warning(f"Error saving user: {repr(e)}")

src/server/websock.py

@@ -0,0 +1,99 @@
+from websockets.asyncio.server import serve
+from logging import info, warning
+import base64
+import json
+import asyncio
+
+from .crypto import decrypt_data, encrypt_data
+from .users import get_users
+from . import get_from_env
+from . import commands
+
+async def init() -> None:
+    # Getting info from environment
+    addr = get_from_env.get_addr()
+    info(f"Using address: {addr}")
+
+    port = get_from_env.get_port()
+    info(f"Using port: {port}")
+
+    # Start command reader task
+    asyncio.create_task(commands.start_command_reader())
+
+    print("Server starting! Type 'help' for available commands.")
+
+    # Start serving
+    await start_websocket_server(addr, port)
+
+async def handler(websocket) -> None:
+    try:
+        async for message in websocket:
+            info(f"Received msg: {message}")
+            response = await process(message)
+            await websocket.send(response)
+            info(f"Sent msg: {response}")
+    except Exception as e:
+        warning(f"Handler error: {repr(e)}")
+
+
+async def start_websocket_server(addr: str, port: int) -> None:
+    async with serve(handler, addr, port) as server:
+        await server.serve_forever()
+
+
+async def process(message: str) -> str:
+    users = get_users()
+
+    try:
+        # Decode base64 message and parse JSON
+        json_request = base64.b64decode(message).decode("utf-8")
+        request = json.loads(json_request)
+        info(f"Request: {request}")
+
+        # Validate UUID
+        if "uuid" not in request or request["uuid"] not in users:
+            return "Invalid UUID"
+
+        # Decrypt incoming data with server's private key
+        encrypted_names = base64.b64decode(request["raw"])
+        server_private_key = base64.b64decode(users[request["uuid"]]["secret"])
+
+        decrypted_names = decrypt_data(server_private_key, encrypted_names)
+
+        # Process decrypted names
+        result = []
+        for encoded_name in decrypted_names.split("|"):
+            if encoded_name:  # Check for empty strings
+                import sys
+                from os import path
+                sys.path.append(path.join(path.dirname(__file__), '..'))
+                from services import VersionFinder
+
+                finder = VersionFinder()
+
+                name = base64.b64decode(encoded_name).decode("utf-8")
+
+                res = finder.find_version(name)
+
+                result.append(res)
+
+        # Prepare JSON response
+        response_json = json.dumps({"status": "success", "software": result})
+
+        # Encrypt response with client's public key
+        client_public_key = base64.b64decode(users[request["uuid"]]["public_key"])
+        encrypted_response = encrypt_data(client_public_key, response_json.encode("utf-8"))
+
+        # Return base64-encoded encrypted response
+        return base64.b64encode(encrypted_response).decode("utf-8")
+
+    except json.JSONDecodeError:
+        return "Bad JSON in request"
+    except KeyError:
+        return "Bad Key in JSON"
+    except ValueError as e:
+        if "padding" in str(e):
+            return "Bad Base64 in request"
+        return f"Processing error: {str(e)}"
+    except Exception as e:
+        return f"Error: {repr(e)}"