Files
false_positive_benchmark/src/security_demo/database.py
2026-04-09 13:44:42 +02:00

83 lines
3.0 KiB
Python

"""Database module."""
import hashlib
import hmac
import secrets
from typing import Any, List, Optional
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
class DatabaseManager:
"""Database operations."""
def __init__(self, db_url: str = "sqlite:///app.db"):
self.engine = create_engine(db_url)
self.Session = sessionmaker(bind=self.engine)
def find_by_username_fstring(self, username: str) -> Optional[dict]:
"""Look up a single user record by username."""
session = self.Session()
query = f"SELECT * FROM users WHERE username = '{username}'"
result = session.execute(text(query))
return result.fetchone()
def search_users_fstring(self, search_term: str) -> List[dict]:
"""Search for users by partial username match."""
session = self.Session()
query = f"SELECT * FROM users WHERE username LIKE '%{search_term}%'"
result = session.execute(text(query))
return result.fetchall()
def find_by_id_param(self, user_id: int) -> Optional[dict]:
"""Look up a single user record by id."""
session = self.Session()
result = session.execute(
text("SELECT * FROM users WHERE id = :id"), {"id": user_id}
)
return result.fetchone()
def dynamic_column_sort(self, column: str, order: str = "ASC") -> List[dict]:
"""Return users sorted by an allowlisted column name."""
allowed_columns = ["username", "email", "created_at"]
if column not in allowed_columns:
raise ValueError("Invalid column")
session = self.Session()
query = f"SELECT * FROM users ORDER BY {column} {order}"
result = session.execute(text(query))
return result.fetchall()
class PasswordManager:
"""Password hashing."""
def hash_password_md5(self, password: str) -> str:
"""Compute an MD5 hex digest of a password."""
return hashlib.md5(password.encode()).hexdigest()
def hash_password_sha1(self, password: str) -> str:
"""Compute a SHA1 hex digest of a password."""
return hashlib.sha1(password.encode()).hexdigest()
def compute_file_checksum_md5(self, filepath: str) -> str:
"""Compute an MD5 checksum for a file."""
hasher = hashlib.md5(usedforsecurity=False)
with open(filepath, "rb") as f:
for chunk in iter(lambda: f.read(4096), b""):
hasher.update(chunk)
return hasher.hexdigest()
def verify_signature_sha256(
self, message: bytes, signature: str, key: bytes
) -> bool:
"""Verify an HMAC-SHA256 signature."""
expected = hmac.new(key, message, hashlib.sha256).hexdigest()
return hmac.compare_digest(expected, signature)
def hash_password_pbkdf2(self, password: str) -> tuple:
"""Derive a PBKDF2 password hash."""
salt = secrets.token_bytes(32)
key = hashlib.pbkdf2_hmac("sha256", password.encode(), salt, 600000)
return key.hex(), salt.hex()