M-CTF-2025/neuralink/build/service/security_service.py

"""
Python implementation of the security service pseudocode.
This module provides a simple CLI-like flow for registering,
logging in, restoring accounts, and managing credentials.
"""

import hashlib
import hmac
import secrets
import sqlite3
import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple


DEFAULT_USERNAME = "default"
USERNAME_MAX_LENGTH = 64
PASSWORD_MAX_LENGTH = 128
IMPLANT_NAME_MAX_LENGTH = 64
IMPLANT_INFO_MAX_LENGTH = 256
SECURITY_CODE_LENGTH = 6
PASSWORD_SALT_BYTES = 16
PBKDF2_ITERATIONS = 200_000
RESTORE_ATTEMPT_LIMIT = 5
RESTORE_WINDOW_SECONDS = 60
RESTORE_LOCK_SECONDS = 30


@dataclass
class SecurityService:
    """Service for managing user credentials in a SQLite database."""

    db_path: str = "security.db"
    username: str = field(default=DEFAULT_USERNAME, init=False)

    def __post_init__(self) -> None:
        self._conn = sqlite3.connect(self.db_path)
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                security_code TEXT NOT NULL
            )
            """
        )
        self._conn.execute(
            """
            CREATE TABLE IF NOT EXISTS implants (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT UNIQUE NOT NULL,
                info TEXT NOT NULL,
                ownername TEXT NOT NULL
            )
            """
        )
        self._conn.commit()
        self._restore_attempts: Dict[str, Dict[str, float]] = {}

    # Utility helpers -----------------------------------------------------
    @staticmethod
    def _generate_random_char() -> str:
        return secrets.choice("abcdefghijklmnopqrstuvwxyz")

    @staticmethod
    def _generate_random_digit() -> str:
        return secrets.choice("0123456789")

    @staticmethod
    def make_unique_username(base_username: str) -> str:
        suffix = "".join(SecurityService._generate_random_char() for _ in range(8))
        return f"{base_username}_{suffix}"

    @staticmethod
    def make_new_password() -> str:
        return "".join(SecurityService._generate_random_char() for _ in range(8))

    @staticmethod
    def generate_code() -> str:
        return "".join(
            SecurityService._generate_random_digit() for _ in range(SECURITY_CODE_LENGTH)
        )

    @staticmethod
    def hash_password(password: str) -> str:
        salt = secrets.token_bytes(PASSWORD_SALT_BYTES)
        digest = hashlib.pbkdf2_hmac(
            "sha256", password.encode("utf-8"), salt, PBKDF2_ITERATIONS
        )
        return f"{salt.hex()}${digest.hex()}"

    @staticmethod
    def _legacy_sha256(password: str) -> str:
        return hashlib.sha256(password.encode("utf-8")).hexdigest()

    @staticmethod
    def is_modern_hash(value: str) -> bool:
        return "$" in value and len(value.split("$", 1)[0]) == PASSWORD_SALT_BYTES * 2

    @staticmethod
    def verify_password(stored_value: str, password: str) -> bool:
        if SecurityService.is_modern_hash(stored_value):
            salt_hex, hash_hex = stored_value.split("$", 1)
            salt = bytes.fromhex(salt_hex)
            candidate = hashlib.pbkdf2_hmac(
                "sha256", password.encode("utf-8"), salt, PBKDF2_ITERATIONS
            )
            return hmac.compare_digest(candidate.hex(), hash_hex)

        legacy_sha = SecurityService._legacy_sha256(password)
        if hmac.compare_digest(stored_value, legacy_sha):
            return True
        return hmac.compare_digest(stored_value, password)

    @staticmethod
    def print_system_data(message: str) -> None:
        print(message)

    @staticmethod
    def _validate_length(value: str, maximum: int, field_name: str) -> Optional[str]:
        trimmed = value.strip()
        if not trimmed:
            print(f"{field_name} cannot be empty.")
            return None
        if len(trimmed) > maximum:
            print(f"{field_name} is too long (max {maximum} characters).")
            return None
        return trimmed

    # Database operations -------------------------------------------------
    def _user_exists(self, username: str) -> bool:
        row = self._conn.execute(
            "SELECT 1 FROM users WHERE username = ?", (username,)
        ).fetchone()
        return row is not None

    def _create_user(self, username: str, password: str, security_code: str) -> bool:
        password_hash = self.hash_password(password)
        try:
            self._conn.execute(
                "INSERT INTO users (username, password_hash, security_code) VALUES (?, ?, ?)",
                (username, password_hash, security_code),
            )
            self._conn.commit()
        except sqlite3.IntegrityError:
            return False
        return True

    def _set_password_hash(self, username: str, password_hash: str) -> None:
        self._conn.execute(
            "UPDATE users SET password_hash = ? WHERE username = ?", (password_hash, username)
        )
        self._conn.commit()

    def _check_user(self, username: str, password: str) -> bool:
        row = self._conn.execute(
            "SELECT password_hash FROM users WHERE username = ?", (username,)
        ).fetchone()
        if row is None:
            return False
        stored_hash = row[0]
        if self.verify_password(stored_hash, password):
            if not self.is_modern_hash(stored_hash):
                self._set_password_hash(username, self.hash_password(password))
            return True
        return False

    def _check_restore_user(self, username: str, code: str) -> bool:
        row = self._conn.execute(
            "SELECT 1 FROM users WHERE username = ? AND security_code = ?",
            (username, code),
        ).fetchone()
        return row is not None

    def _change_password(self, username: str, new_password: str) -> bool:
        password_hash = self.hash_password(new_password)
        cursor = self._conn.execute(
            "UPDATE users SET password_hash = ? WHERE username = ?",
            (password_hash, username),
        )
        self._conn.commit()
        return cursor.rowcount > 0

    def _get_security_code(self, username: str) -> Optional[str]:
        row = self._conn.execute(
            "SELECT security_code FROM users WHERE username = ?", (username,)
        ).fetchone()
        return row[0] if row else None

    def _user_implants(self, username: str) -> List[str]:
        rows = self._conn.execute(
            "SELECT name FROM implants WHERE ownername = ?", (username,)
        ).fetchall()
        return [row[0] for row in rows]

    def _implant_belongs_to(self, username: str, name: str) -> bool:
        row = self._conn.execute(
            "SELECT 1 FROM implants WHERE ownername = ? AND name = ?",
            (username, name),
        ).fetchone()
        return row is not None

    def _delete_implant(self, name: str) -> bool:
        cursor = self._conn.execute("DELETE FROM implants WHERE name = ?", (name,))
        self._conn.commit()
        return cursor.rowcount > 0

    def _add_implant(self, name: str, info: str, ownername: str) -> bool:
        try:
            self._conn.execute(
                "INSERT INTO implants (name, info, ownername) VALUES (?, ?, ?)",
                (name, info, ownername),
            )
            self._conn.commit()
        except sqlite3.IntegrityError:
            return False
        return True

    def _get_implant_info(self, name: str) -> Optional[str]:
        row = self._conn.execute(
            "SELECT info FROM implants WHERE name = ?", (name,)
        ).fetchone()
        return row[0] if row else None

    # User flows ----------------------------------------------------------
    def register_user(self) -> None:
        raw_username = input("\nEnter username: ")
        username = self._validate_length(raw_username, USERNAME_MAX_LENGTH, "Username")
        if username is None:
            return
        password_input = input("Enter password: ")
        password = self._validate_length(
            password_input, PASSWORD_MAX_LENGTH, "Password"
        )
        if password is None:
            return
        self.print_system_data("Creating new account. Please wait...")

        candidate_username = self.make_unique_username(username)
        security_code = self.generate_code()

        if self._user_exists(candidate_username):
            self.print_system_data("User already exists.")
            return

        if self._create_user(candidate_username, password, security_code):
            self.print_system_data("---Your credentials ---")
            print(f"Username: {candidate_username}")
            print(f"Password: {password}")
            print(f"Security code: {security_code}")
            self.print_system_data("Use these credentials to gain access to the system.")
        else:
            self.print_system_data("Failed to create user.")

    def login_user(self) -> None:
        username_input = input("\nEnter username: ")
        username = self._validate_length(
            username_input, USERNAME_MAX_LENGTH, "Username"
        )
        if username is None:
            return
        password_input = input("Enter password: ")
        password = self._validate_length(
            password_input, PASSWORD_MAX_LENGTH, "Password"
        )
        if password is None:
            return
        self.print_system_data("Trying to log in...")

        if self._check_user(username, password):
            self.username = username
            self.print_system_data("Successfully logged in.")
        else:
            self.print_system_data("Failed to log in.")

    def restore_user(self) -> None:
        username_input = input("\nEnter username: ")
        username = self._validate_length(
            username_input, USERNAME_MAX_LENGTH, "Username"
        )
        if username is None:
            return
        allowed, wait_time = self._can_attempt_restore(username)
        if not allowed:
            self.print_system_data(
                f"Too many attempts. Try again in {int(wait_time) + 1} seconds."
            )
            return

        code = input("Enter security code: ").strip()
        self.print_system_data("Trying to find user...")

        if not self._check_restore_user(username, code):
            self._record_failed_restore(username)
            self.print_system_data("Failed to find user.")
            return

        self._clear_restore_attempts(username)
        self.print_system_data("Successfully found user.")
        new_password = self.make_new_password()
        if self._change_password(username, new_password):
            self.print_system_data("Changing password...")
            self.print_system_data("---Your new credentials ---")
            print(f"Username: {username}")
            print(f"Password: {new_password}")
            self.print_system_data(
                "Use these credentials to gain access to the system."
            )
        else:
            self.print_system_data("Unexpected error. Please try later.")

    def change_password(self) -> None:
        if self.username == DEFAULT_USERNAME:
            self.print_system_data("You need to log in first.")
            return

        new_pass_input = input("Enter a new password: ")
        new_password = self._validate_length(
            new_pass_input, PASSWORD_MAX_LENGTH, "Password"
        )
        if new_password is None:
            return

        if self._change_password(self.username, new_password):
            self.print_system_data("Password changed successfully.")
        else:
            self.print_system_data("Failed to change password.")

    def show_security_code(self) -> None:
        if self.username == DEFAULT_USERNAME:
            self.print_system_data("You need to log in first.")
            return

        code = self._get_security_code(self.username)
        if code is None:
            self.print_system_data("Failed to retrieve security code.")
        else:
            print(f"Security code: {code}")

    def add_implant(self) -> None:
        if self.username == DEFAULT_USERNAME:
            self.print_system_data("You need to log in first.")
            return

        name_input = input("\nEnter implant name: ")
        name = self._validate_length(
            name_input, IMPLANT_NAME_MAX_LENGTH, "Implant name"
        )
        if name is None:
            self.print_system_data("Invalid option selected.")
            return
        info_input = input("\nEnter implant info: ")
        info = self._validate_length(
            info_input, IMPLANT_INFO_MAX_LENGTH, "Implant info"
        )
        if info is None:
            self.print_system_data("Invalid option selected.")
            return

        if self._add_implant(name, info, self.username):
            self.print_system_data("Implant added successfully.")
        else:
            self.print_system_data("Unexpected error. Please try later.")

    def delete_implant(self) -> None:
        if self.username == DEFAULT_USERNAME:
            self.print_system_data("You need to log in first.")
            return

        self.print_system_data("Getting list of implants. Please wait...")
        implants = self._user_implants(self.username)
        if not implants:
            self.print_system_data("No implants found for this user.")
            return

        for idx, implant in enumerate(implants, start=1):
            print(f"{idx}. {implant}")

        name_input = input("\nEnter implant name: ")
        name = self._validate_length(
            name_input, IMPLANT_NAME_MAX_LENGTH, "Implant name"
        )
        if name is None:
            self.print_system_data("Invalid option selected.")
            return

        if not self._implant_belongs_to(self.username, name):
            self.print_system_data("Invalid option selected.")
            return

        self.print_system_data("Deleting implant...")
        if self._delete_implant(name):
            self.print_system_data("Implant deleted successfully.")
        else:
            self.print_system_data("Unexpected error. Please try later.")

    def show_implant_info(self) -> None:
        if self.username == DEFAULT_USERNAME:
            self.print_system_data("You need to log in first.")
            return

        self.print_system_data("Getting list of implants. Please wait...")
        implants = self._user_implants(self.username)
        if not implants:
            self.print_system_data("No implants found for this user.")
            return

        for idx, implant in enumerate(implants, start=1):
            print(f"{idx}. {implant}")

        name_input = input("\nEnter implant name: ")
        name = self._validate_length(
            name_input, IMPLANT_NAME_MAX_LENGTH, "Implant name"
        )
        if name is None:
            self.print_system_data("Invalid option selected.")
            return

        if not self._implant_belongs_to(self.username, name):
            self.print_system_data("Invalid option selected.")
            return

        self.print_system_data("Getting implant info...")
        info = self._get_implant_info(name)
        if info is None:
            self.print_system_data("Unexpected error. Please try later.")
            return

        print(f"Implant info: {info}")

    # Menus ---------------------------------------------------------------
    def settings_menu(self) -> None:
        while self.username != DEFAULT_USERNAME:
            self.print_system_data("--- Settings menu ---")
            self.print_system_data("1. Change password")
            self.print_system_data("2. Show security code")
            self.print_system_data("3. Return to previous menu")
            choice = input("Choose an option (1-3): ").strip()
            if choice == "1":
                self.change_password()
            elif choice == "2":
                self.show_security_code()
            elif choice == "3":
                return
            else:
                self.print_system_data("Invalid option selected.")

    def app_menu(self) -> None:
        while self.username != DEFAULT_USERNAME:
            self.print_system_data("--- Application menu ---")
            self.print_system_data("1. Implant menu")
            self.print_system_data("2. Settings menu")
            self.print_system_data("3. Log out")
            choice = input("Choose an option (1-3): ").strip()
            if choice == "1":
                self.implants_menu()
            elif choice == "2":
                self.settings_menu()
            elif choice == "3":
                self.username = DEFAULT_USERNAME
                self.print_system_data("Logged out.")
            else:
                self.print_system_data("Invalid option selected.")

    def implants_menu(self) -> None:
        while self.username != DEFAULT_USERNAME:
            self.print_system_data("--- Implants menu ---")
            self.print_system_data("1. Add new implant")
            self.print_system_data("2. Delete implant")
            self.print_system_data("3. Show info about implant")
            self.print_system_data("4. Return to previous menu")
            choice = input("Choose an option (1-4): ").strip()
            if choice == "1":
                self.add_implant()
            elif choice == "2":
                self.delete_implant()
            elif choice == "3":
                self.show_implant_info()
            elif choice == "4":
                return
            else:
                self.print_system_data("Invalid option selected.")

    def startup_menu(self) -> None:
        while True:
            if self.username != DEFAULT_USERNAME:
                self.app_menu()
                continue

            self.print_system_data("--- Startup menu ---")
            self.print_system_data("1. Register new account")
            self.print_system_data("2. Login to account")
            self.print_system_data("3. Restore an account")
            self.print_system_data("4. Exit from program")
            choice = input("Choose an option (1-4): ").strip()
            if choice == "1":
                self.register_user()
            elif choice == "2":
                self.login_user()
            elif choice == "3":
                self.restore_user()
            elif choice == "4":
                self.print_system_data("Goodbye!")
                break
            else:
                self.print_system_data("Invalid option selected.")

    # Rate limiting -------------------------------------------------------
    def _can_attempt_restore(self, username: str) -> Tuple[bool, float]:
        now = time.time()
        stats = self._restore_attempts.get(username)
        if not stats:
            return True, 0.0
        locked_until = stats.get("locked_until", 0.0)
        if locked_until > now:
            return False, locked_until - now
        window_start = stats.get("window_start", now)
        if now - window_start > RESTORE_WINDOW_SECONDS:
            self._restore_attempts.pop(username, None)
            return True, 0.0
        return True, 0.0

    def _record_failed_restore(self, username: str) -> None:
        now = time.time()
        stats = self._restore_attempts.get(username)
        if not stats or now - stats.get("window_start", now) > RESTORE_WINDOW_SECONDS:
            stats = {"count": 0, "window_start": now, "locked_until": 0.0}
        stats["count"] += 1
        stats["window_start"] = stats.get("window_start", now)
        if stats["count"] >= RESTORE_ATTEMPT_LIMIT:
            stats["locked_until"] = now + RESTORE_LOCK_SECONDS
            stats["count"] = 0
            stats["window_start"] = stats["locked_until"]
        self._restore_attempts[username] = stats

    def _clear_restore_attempts(self, username: str) -> None:
        self._restore_attempts.pop(username, None)


def main() -> None:
    service = SecurityService()
    service.print_system_data("Configuring network interfaces... done")
    service.print_system_data("Mounting /dev/sda1... done")
    service.print_system_data("Starting random number generator daemon... done")
    service.startup_menu()


if __name__ == "__main__":
    main()