Download database/database_layer.py from samliumay/local-llm-pentesting-lab: direct link, hf CLI and curl.
- Browser
- Download file 17.9 kB
-
https://huggingface.co/spaces/samliumay/local-llm-pentesting-lab/resolve/main/database/database_layer.py
- Command line
-
hf download hf://spaces/samliumay/local-llm-pentesting-lab/database/database_layer.py
-
curl -L -o database_layer.py https://huggingface.co/spaces/samliumay/local-llm-pentesting-lab/resolve/main/database/database_layer.py
17.9 kB
| """PostgreSQL access methods used by authentication and model tools. | |
| Trusted application code should prefer narrow, parameterized methods such as | |
| ``get_customer_account_summary``. ``execute_query`` remains an intentionally raw | |
| lab primitive and is not exposed to the Ollama model. | |
| """ | |
| import os | |
| from decimal import Decimal, InvalidOperation | |
| from pathlib import Path | |
| from typing import Any | |
| from uuid import UUID | |
| import psycopg | |
| from dotenv import load_dotenv | |
| class database_layer: | |
| """Own database configuration and the repository's PostgreSQL operations.""" | |
| def __init__(self): | |
| """Load local connection settings from ``configuration/.env``.""" | |
| self.database_env_file = Path(__file__).resolve().parents[1] / "configuration" / ".env" | |
| # Explicitly load the repository-local file instead of relying on cwd. | |
| load_dotenv(self.database_env_file) | |
| # Defaults make the lab easy to start; deployments should override them. | |
| self._connection_settings = { | |
| "host": os.getenv("POSTGRES_HOST", "127.0.0.1"), | |
| "port": os.getenv("POSTGRES_PORT", "5432"), | |
| "dbname": os.getenv("POSTGRES_DB", "world_order_bank_system_db"), | |
| "user": os.getenv("POSTGRES_USER", "username"), | |
| "password": os.getenv("POSTGRES_PASSWORD", "1234") | |
| } | |
| # dotenv values are strings; currently even "False" is truthy here. The two | |
| # execution paths are equivalent until firewall parsing and policy are implemented. | |
| self._firewall_controls = os.getenv("firewall_controls") | |
| print("Be sure that you configured your enviorment files correctly.") | |
| def verify_customer_credentials(self, email: str, password: str) -> bool: | |
| """Check a customer's password without exposing its stored hash.""" | |
| # Parameters are bound separately so credentials never become SQL text. | |
| query = """ | |
| SELECT EXISTS ( | |
| SELECT 1 | |
| FROM customers | |
| WHERE email = %s | |
| AND password_hash = crypt(%s, password_hash) | |
| ) | |
| """ | |
| with psycopg.connect(**self._connection_settings) as connection: | |
| with connection.cursor() as cursor: | |
| cursor.execute(query, (email, password)) | |
| result = cursor.fetchone() | |
| return bool(result and result[0]) | |
| def get_customer_account_summary(self, email: str) -> dict[str, Any] | None: | |
| """Return permitted summary fields for exactly one authenticated email. | |
| The caller—not the model—supplies ``email``. Returning ``None`` means that | |
| no matching customer exists. | |
| """ | |
| if not isinstance(email, str) or not email.strip(): | |
| raise ValueError("The customer email cannot be empty.") | |
| query = """ | |
| SELECT first_name, last_name, email, total_amount, currency_code | |
| FROM customers | |
| WHERE email = %s | |
| """ | |
| with psycopg.connect(**self._connection_settings) as connection: | |
| with connection.cursor() as cursor: | |
| cursor.execute(query, (email.strip(),)) | |
| result = cursor.fetchone() | |
| if result is None: | |
| return None | |
| return { | |
| "first_name": result[0], | |
| "last_name": result[1], | |
| "email": result[2], | |
| "total_amount": result[3], | |
| "currency_code": result[4], | |
| } | |
| # Convert model- or application-supplied amounts into exact PostgreSQL money | |
| # values without rounding, accepting booleans, or carrying float artifacts. | |
| def _normalise_transfer_amount(amount: Decimal | str | int | float) -> Decimal: | |
| """Return a positive two-decimal amount suitable for ``NUMERIC(19, 2)``.""" | |
| # bool is an int subclass in Python, but True must never become a monetary 1. | |
| if isinstance(amount, bool): | |
| raise ValueError("The transfer amount must be a number.") | |
| try: | |
| # Convert through text so JSON numbers do not carry binary-float noise into SQL. | |
| decimal_amount = Decimal(str(amount)) | |
| except (InvalidOperation, ValueError): | |
| raise ValueError("The transfer amount must be a valid decimal number.") from None | |
| if not decimal_amount.is_finite() or decimal_amount <= 0: | |
| raise ValueError("The transfer amount must be greater than zero.") | |
| try: | |
| # Quantize for validation only; silently rounding money would change user intent. | |
| normalised_amount = decimal_amount.quantize(Decimal("0.01")) | |
| except InvalidOperation: | |
| raise ValueError("The transfer amount is outside the supported range.") from None | |
| if normalised_amount != decimal_amount: | |
| raise ValueError("The transfer amount cannot have more than two decimal places.") | |
| if normalised_amount > Decimal("99999999999999999.99"): | |
| raise ValueError("The transfer amount is outside the supported range.") | |
| return normalised_amount | |
| # Validate a proposed transfer and create its pending ledger record. The sender | |
| # is trusted application state, and this phase never changes customer balances. | |
| def prepare_transfer( | |
| self, | |
| authenticated_sender_email: str, | |
| recipient_email: str, | |
| amount: Decimal | str | int | float, | |
| ) -> dict[str, Any]: | |
| """Create a pending transfer without changing either account balance.""" | |
| if not isinstance(authenticated_sender_email, str) or not authenticated_sender_email.strip(): | |
| raise ValueError("The authenticated sender email cannot be empty.") | |
| if not isinstance(recipient_email, str) or not recipient_email.strip(): | |
| raise ValueError("The recipient email cannot be empty.") | |
| sender_email = authenticated_sender_email.strip() | |
| recipient_email = recipient_email.strip() | |
| # The sender originates in authenticated application state, never in tool arguments. | |
| if sender_email == recipient_email: | |
| raise ValueError("The sender and recipient must be different customers.") | |
| transfer_amount = self._normalise_transfer_amount(amount) | |
| customer_query = """ | |
| SELECT customer_id, first_name, last_name, email, total_amount, currency_code | |
| FROM customers | |
| WHERE email IN (%s, %s) | |
| """ | |
| insert_query = """ | |
| INSERT INTO transfers ( | |
| sender_customer_id, | |
| recipient_customer_id, | |
| amount, | |
| currency_code | |
| ) | |
| VALUES (%s, %s, %s, %s) | |
| RETURNING confirmation_id, status, created_at, expires_at | |
| """ | |
| with psycopg.connect(**self._connection_settings) as connection: | |
| with connection.cursor() as cursor: | |
| # Preparation reads both accounts but deliberately does not lock or update them. | |
| # Confirmation repeats every mutable check under row locks. | |
| cursor.execute(customer_query, (sender_email, recipient_email)) | |
| customers = {row[3]: row for row in cursor.fetchall()} | |
| sender = customers.get(sender_email) | |
| if sender is None: | |
| raise ValueError("The authenticated sender account was not found.") | |
| recipient = customers.get(recipient_email) | |
| if recipient is None: | |
| raise ValueError("The recipient account was not found.") | |
| sender_currency = str(sender[5]).strip() | |
| recipient_currency = str(recipient[5]).strip() | |
| # Moving a numeric amount between currencies would invent an exchange rate. | |
| if sender_currency != recipient_currency: | |
| raise ValueError("Transfers between different currencies are not supported.") | |
| if sender[4] < transfer_amount: | |
| raise ValueError("The account has insufficient funds for this transfer.") | |
| cursor.execute( | |
| insert_query, | |
| (sender[0], recipient[0], transfer_amount, sender_currency), | |
| ) | |
| # This insert is the only write during preparation; balances are untouched. | |
| transfer = cursor.fetchone() | |
| if transfer is None: | |
| raise RuntimeError("The pending transfer could not be created.") | |
| return { | |
| "confirmation_id": str(transfer[0]), | |
| "status": transfer[1], | |
| "sender_email": sender_email, | |
| "recipient_email": recipient_email, | |
| "recipient_name": f"{recipient[1]} {recipient[2]}", | |
| "amount": transfer_amount, | |
| "currency_code": sender_currency, | |
| "created_at": transfer[2], | |
| "expires_at": transfer[3], | |
| } | |
| # Complete one sender-owned pending transfer inside a single transaction. Row | |
| # locks and repeated validation prevent replay and stale-balance updates. | |
| def confirm_transfer( | |
| self, | |
| authenticated_sender_email: str, | |
| confirmation_id: str, | |
| ) -> dict[str, Any]: | |
| """Atomically complete a pending transfer owned by the authenticated sender.""" | |
| if not isinstance(authenticated_sender_email, str) or not authenticated_sender_email.strip(): | |
| raise ValueError("The authenticated sender email cannot be empty.") | |
| if not isinstance(confirmation_id, str) or not confirmation_id.strip(): | |
| raise ValueError("The confirmation ID cannot be empty.") | |
| sender_email = authenticated_sender_email.strip() | |
| try: | |
| # Parse before querying so malformed identifiers never reach PostgreSQL. | |
| parsed_confirmation_id = UUID(confirmation_id.strip()) | |
| except ValueError: | |
| raise ValueError("The confirmation ID is not valid.") from None | |
| transfer_query = """ | |
| SELECT | |
| t.transfer_id, | |
| t.status, | |
| t.amount, | |
| t.currency_code, | |
| t.sender_customer_id, | |
| t.recipient_customer_id, | |
| sender.email, | |
| recipient.email, | |
| t.created_at, | |
| t.expires_at, | |
| t.confirmed_at | |
| FROM transfers AS t | |
| JOIN customers AS sender | |
| ON sender.customer_id = t.sender_customer_id | |
| JOIN customers AS recipient | |
| ON recipient.customer_id = t.recipient_customer_id | |
| WHERE t.confirmation_id = %s | |
| AND sender.email = %s | |
| FOR UPDATE OF t | |
| """ | |
| lock_customers_query = """ | |
| SELECT customer_id, total_amount, currency_code | |
| FROM customers | |
| WHERE customer_id IN (%s, %s) | |
| ORDER BY customer_id | |
| FOR UPDATE | |
| """ | |
| with psycopg.connect(**self._connection_settings) as connection: | |
| with connection.cursor() as cursor: | |
| # Locking the ledger row serializes concurrent confirmation and replay attempts. | |
| # Scoping by sender email prevents one customer confirming another's transfer. | |
| cursor.execute(transfer_query, (parsed_confirmation_id, sender_email)) | |
| transfer = cursor.fetchone() | |
| if transfer is None: | |
| raise ValueError( | |
| "The transfer was not found for the authenticated customer." | |
| ) | |
| transfer_id = transfer[0] | |
| status = transfer[1] | |
| transfer_amount = transfer[2] | |
| transfer_currency = str(transfer[3]).strip() | |
| sender_customer_id = transfer[4] | |
| recipient_customer_id = transfer[5] | |
| recipient_email = transfer[7] | |
| if status == "completed": | |
| # Replays report the previous outcome without moving money again. | |
| return { | |
| "confirmation_id": str(parsed_confirmation_id), | |
| "status": "already_completed", | |
| "sender_email": sender_email, | |
| "recipient_email": recipient_email, | |
| "amount": transfer_amount, | |
| "currency_code": transfer_currency, | |
| "confirmed_at": transfer[10], | |
| } | |
| if status == "expired": | |
| return { | |
| "confirmation_id": str(parsed_confirmation_id), | |
| "status": "expired", | |
| } | |
| cursor.execute( | |
| """ | |
| UPDATE transfers | |
| SET status = 'expired' | |
| WHERE transfer_id = %s | |
| AND status = 'pending' | |
| AND expires_at <= CURRENT_TIMESTAMP | |
| RETURNING status | |
| """, | |
| (transfer_id,), | |
| ) | |
| if cursor.fetchone() is not None: | |
| # Persist the terminal expired state while leaving balances unchanged. | |
| return { | |
| "confirmation_id": str(parsed_confirmation_id), | |
| "status": "expired", | |
| } | |
| cursor.execute( | |
| lock_customers_query, | |
| (sender_customer_id, recipient_customer_id), | |
| ) | |
| # Stable customer-ID ordering prevents opposite-direction transfers from | |
| # taking the same pair of account locks in conflicting orders. | |
| customer_rows = cursor.fetchall() | |
| customers = {row[0]: row for row in customer_rows} | |
| sender = customers.get(sender_customer_id) | |
| recipient = customers.get(recipient_customer_id) | |
| if sender is None or recipient is None: | |
| raise RuntimeError("A transfer account no longer exists.") | |
| sender_currency = str(sender[2]).strip() | |
| recipient_currency = str(recipient[2]).strip() | |
| # Balance and currency may have changed since preparation, so check again. | |
| if ( | |
| sender_currency != transfer_currency | |
| or recipient_currency != transfer_currency | |
| ): | |
| raise ValueError("The account currencies no longer match the transfer.") | |
| if sender[1] < transfer_amount: | |
| raise ValueError("The account has insufficient funds for this transfer.") | |
| cursor.execute( | |
| """ | |
| UPDATE customers | |
| SET total_amount = total_amount - %s, | |
| updated_at = CURRENT_TIMESTAMP | |
| WHERE customer_id = %s | |
| """, | |
| (transfer_amount, sender_customer_id), | |
| ) | |
| cursor.execute( | |
| """ | |
| UPDATE customers | |
| SET total_amount = total_amount + %s, | |
| updated_at = CURRENT_TIMESTAMP | |
| WHERE customer_id = %s | |
| """, | |
| (transfer_amount, recipient_customer_id), | |
| ) | |
| cursor.execute( | |
| """ | |
| UPDATE transfers | |
| SET status = 'completed', | |
| confirmed_at = CURRENT_TIMESTAMP | |
| WHERE transfer_id = %s | |
| AND status = 'pending' | |
| RETURNING confirmed_at | |
| """, | |
| (transfer_id,), | |
| ) | |
| completion = cursor.fetchone() | |
| if completion is None: | |
| raise RuntimeError("The transfer could not be confirmed.") | |
| # The connection context commits all three updates together. Any exception | |
| # rolls back both balance changes and the ledger status change. | |
| return { | |
| "confirmation_id": str(parsed_confirmation_id), | |
| "status": "completed", | |
| "sender_email": sender_email, | |
| "recipient_email": recipient_email, | |
| "amount": transfer_amount, | |
| "currency_code": transfer_currency, | |
| "confirmed_at": completion[0], | |
| "remaining_balance": sender[1] - transfer_amount, | |
| } | |
| def execute_query(self, query:str) -> str: | |
| """Execute raw SQL through the selected lab path. | |
| This intentionally dangerous primitive supports future firewall experiments. | |
| It must not be exposed directly as an LLM tool. | |
| """ | |
| if(self._firewall_controls): | |
| return self._secure_query_execute(query) | |
| return self._non_secure_query_execute(query) | |
| # TODO: enforce an actual SQL policy. This currently behaves like the unsafe path. | |
| def _secure_query_execute(self, query:str ) -> str: | |
| """Execute SQL through the placeholder firewall-controlled path.""" | |
| with psycopg.connect( | |
| **self._connection_settings | |
| ) as connection: | |
| with connection.cursor() as cursor: | |
| cursor.execute(query) | |
| answer = cursor.fetchall() | |
| return answer | |
| def _non_secure_query_execute(self, query:str) -> str: | |
| """Execute unrestricted SQL for controlled pentesting experiments.""" | |
| with psycopg.connect( | |
| **self._connection_settings | |
| ) as connection: | |
| with connection.cursor() as cursor: | |
| cursor.execute(query) | |
| answer = cursor.fetchall() | |
| return answer | |