551 lines
20 KiB
Python
551 lines
20 KiB
Python
"""
|
|
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()
|