265 lines
9.6 KiB
Python
265 lines
9.6 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 secrets
|
||
|
|
import sqlite3
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
|
||
|
|
DEFAULT_USERNAME = "default"
|
||
|
|
|
||
|
|
|
||
|
|
@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.commit()
|
||
|
|
|
||
|
|
# 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(3))
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def hash_password(password: str) -> str:
|
||
|
|
digest = hashlib.sha256(password.encode("utf-8")).hexdigest()
|
||
|
|
return digest
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def print_system_data(message: str) -> None:
|
||
|
|
print(message)
|
||
|
|
|
||
|
|
# 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 _check_user(self, username: str, password: str) -> bool:
|
||
|
|
password_hash = self.hash_password(password)
|
||
|
|
row = self._conn.execute(
|
||
|
|
"SELECT 1 FROM users WHERE username = ? AND password_hash = ?",
|
||
|
|
(username, password_hash),
|
||
|
|
).fetchone()
|
||
|
|
return row is not None
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
# User flows ----------------------------------------------------------
|
||
|
|
def register_user(self) -> None:
|
||
|
|
raw_username = input("\nEnter username: ").strip()
|
||
|
|
password = input("Enter password: ").strip()
|
||
|
|
self.print_system_data("Creating new account. Please wait...")
|
||
|
|
|
||
|
|
candidate_username = self.make_unique_username(raw_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("\nEnter username: ").strip()
|
||
|
|
password = input("Enter password: ").strip()
|
||
|
|
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("\nEnter username: ").strip()
|
||
|
|
code = input("Enter security code: ").strip()
|
||
|
|
self.print_system_data("Trying to find user...")
|
||
|
|
|
||
|
|
if not self._check_restore_user(username, code):
|
||
|
|
self.print_system_data("Failed to find user.")
|
||
|
|
return
|
||
|
|
|
||
|
|
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_password = input("Enter a new password: ").strip()
|
||
|
|
if not new_password:
|
||
|
|
self.print_system_data("Password cannot be empty.")
|
||
|
|
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}")
|
||
|
|
|
||
|
|
# 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. Settings menu")
|
||
|
|
self.print_system_data("2. Log out")
|
||
|
|
choice = input("Choose an option (1-2): ").strip()
|
||
|
|
if choice == "1":
|
||
|
|
self.settings_menu()
|
||
|
|
elif choice == "2":
|
||
|
|
self.username = DEFAULT_USERNAME
|
||
|
|
self.print_system_data("Logged out.")
|
||
|
|
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.")
|
||
|
|
|
||
|
|
|
||
|
|
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()
|