walletpress/backend/core/hosting.py

"""WalletPress Hosted — multi-tenant user management + usage tracking.

Revenue model:
  Free:    3 chains, 10 wallets/day, no API access
  Starter: $29/mo, 10 chains, 500 wallets/day, API access
  Pro:     $79/mo, 55 chains, unlimited wallets, full API, 2FA
  Enterprise: $199/mo, everything + dedicated support + SLA

Users self-register, get API keys, and pay via Stripe.
Admin creates plans, views usage, manages subscriptions.

SECURITY: Passwords are hashed with Argon2id (m=64MB, t=3, p=4).
Legacy SHA-256 hashes are migrated on next successful login.
"""

from __future__ import annotations

import logging
import secrets
import sqlite3
import time
from pathlib import Path
from typing import Optional

from argon2 import PasswordHasher
from argon2.exceptions import (
    VerifyMismatchError,
    InvalidHashError,
    VerificationError,
)

from .config import cfg

logger = logging.getLogger("wp.hosting")

# Argon2id parameters: OWASP recommendations as of 2024
# m = 64 MiB memory cost, t = 3 iterations, p = 4 parallelism
# (argon2-cffi defaults are 64MB / 3 iterations / 4 lanes)
_ph = PasswordHasher()


def _hash_password(plaintext: str) -> str:
    """Hash a password with Argon2id. Returns the encoded hash string.

    The returned string includes the salt, parameters, and hash in a
    self-describing format like:
        $argon2id$v=19$m=65536,t=3,p=4$<salt>$<hash>
    """
    return _ph.hash(plaintext)


def _verify_password(stored_hash: str, plaintext: str) -> bool:
    """Verify a plaintext password against a stored hash.

    Supports both new Argon2id hashes AND legacy unsalted SHA-256 hashes
    (for migration). Returns True on match, False on mismatch.
    """
    if not stored_hash:
        return False
    # Legacy SHA-256 hash: 64 hex chars, no $ separators
    if "$" not in stored_hash and len(stored_hash) == 64:
        try:
            import hashlib
            legacy_match = hashlib.sha256(plaintext.encode()).hexdigest() == stored_hash
            return legacy_match
        except Exception:
            return False
    # Argon2id hash
    try:
        _ph.verify(stored_hash, plaintext)
        return True
    except (VerifyMismatchError, VerificationError, InvalidHashError):
        return False


def _needs_rehash(stored_hash: str) -> bool:
    """Check if the stored hash uses outdated parameters and needs re-hashing."""
    if "$" not in stored_hash:
        return True  # Legacy SHA-256 — needs migration to Argon2id
    try:
        return _ph.check_needs_rehash(stored_hash)
    except InvalidHashError:
        return True


PLANS = {
    "free": {"name": "Free", "price_usd": 0, "chains": 3, "daily_wallet_limit": 10, "api_access": False, "features": ["basic_generation"]},
    "starter": {"name": "Starter", "price_usd": 29, "chains": 10, "daily_wallet_limit": 500, "api_access": True, "features": ["basic_generation", "batch", "mnemonic", "export"]},
    "pro": {"name": "Pro", "price_usd": 79, "chains": 55, "daily_wallet_limit": 0, "api_access": True, "features": ["all"]},
    "enterprise": {"name": "Enterprise", "price_usd": 199, "chains": 55, "daily_wallet_limit": 0, "api_access": True, "features": ["all", "2fa", "sla"]},
}


class HostingDB:
    """User and subscription management."""

    def __init__(self, db_path: Path):
        self.db_path = db_path
        self._init_db()

    def _conn(self):
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        return conn

    def _init_db(self):
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        conn = self._conn()
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                id TEXT PRIMARY KEY,
                email TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                name TEXT DEFAULT '',
                plan TEXT DEFAULT 'free',
                api_key TEXT UNIQUE DEFAULT '',
                stripe_customer_id TEXT DEFAULT '',
                subscription_id TEXT DEFAULT '',
                subscription_status TEXT DEFAULT 'active',
                created_at REAL NOT NULL,
                last_login_at REAL DEFAULT 0
            );
            CREATE TABLE IF NOT EXISTS usage_log (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id TEXT NOT NULL,
                action TEXT NOT NULL,
                chain TEXT DEFAULT '',
                count INTEGER DEFAULT 1,
                timestamp REAL NOT NULL,
                ip TEXT DEFAULT '',
                FOREIGN KEY(user_id) REFERENCES users(id)
            );
            CREATE INDEX IF NOT EXISTS idx_usage_user ON usage_log(user_id);
            CREATE INDEX IF NOT EXISTS idx_usage_date ON usage_log(timestamp);
        """)
        conn.commit()
        conn.close()

    def register(self, email: str, password: str, name: str = "") -> dict:
        conn = self._conn()
        existing = conn.execute("SELECT id FROM users WHERE email = ?", (email,)).fetchone()
        if existing:
            conn.close()
            return {"error": "Email already registered"}
        user_id = f"u_{secrets.token_hex(8)}"
        pw_hash = _hash_password(password)
        api_key = f"wp_live_{secrets.token_hex(24)}"
        conn.execute(
            "INSERT INTO users (id, email, password_hash, name, plan, api_key, created_at) VALUES (?, ?, ?, ?, 'free', ?, ?)",
            (user_id, email, pw_hash, name, api_key, time.time()),
        )
        conn.commit()
        conn.close()
        return {"user_id": user_id, "api_key": api_key, "plan": "free"}

    def login(self, email: str, password: str) -> Optional[dict]:
        """Verify login and migrate legacy SHA-256 hash to Argon2id on success.

        Returns the user row (without password_hash) on success, None on failure.
        On successful login with an outdated hash, transparently re-hashes the
        password with current Argon2id parameters.
        """
        conn = self._conn()
        row = conn.execute("SELECT * FROM users WHERE email = ?", (email,)).fetchone()
        if not row:
            conn.close()
            return None
        stored_hash = row["password_hash"]
        if not _verify_password(stored_hash, password):
            conn.close()
            return None

        # Migrate outdated hash to current Argon2id params
        if _needs_rehash(stored_hash):
            try:
                new_hash = _hash_password(password)
                conn.execute(
                    "UPDATE users SET password_hash = ? WHERE id = ?",
                    (new_hash, row["id"]),
                )
                conn.commit()
                logger.info(f"Migrated password hash to Argon2id for user {row['id']}")
            except Exception as e:
                logger.warning(f"Password hash migration failed for user {row['id']}: {e}")

        # Return everything except password_hash — never leak it
        safe = {k: v for k, v in dict(row).items() if k != "password_hash"}
        conn.close()
        return safe

    def get_by_api_key(self, api_key: str) -> Optional[dict]:
        conn = self._conn()
        row = conn.execute("SELECT * FROM users WHERE api_key = ?", (api_key,)).fetchone()
        conn.close()
        return dict(row) if row else None

    def log_usage(self, user_id: str, action: str, chain: str = "", count: int = 1, ip: str = ""):
        conn = self._conn()
        conn.execute(
            "INSERT INTO usage_log (user_id, action, chain, count, timestamp, ip) VALUES (?, ?, ?, ?, ?, ?)",
            (user_id, action, chain, count, time.time(), ip),
        )
        conn.commit()
        conn.close()

    def daily_usage(self, user_id: str) -> int:
        today = int(time.time()) - (int(time.time()) % 86400)
        conn = self._conn()
        row = conn.execute(
            "SELECT COALESCE(SUM(count), 0) as total FROM usage_log WHERE user_id = ? AND timestamp >= ?",
            (user_id, today),
        ).fetchone()
        conn.close()
        return row["total"] if row else 0

    def check_limit(self, user_id: str) -> tuple[bool, str]:
        conn = self._conn()
        row = conn.execute("SELECT plan FROM users WHERE id = ?", (user_id,)).fetchone()
        conn.close()
        if not row:
            return False, "User not found"
        plan = PLANS.get(row["plan"], PLANS["free"])
        if plan["daily_wallet_limit"] == 0:
            return True, ""  # Unlimited
        used = self.daily_usage(user_id)
        if used >= plan["daily_wallet_limit"]:
            return False, f"Daily limit reached ({plan['daily_wallet_limit']} wallets). Upgrade at /hosting/plans"
        return True, ""

    def get_user(self, user_id: str) -> Optional[dict]:
        conn = self._conn()
        row = conn.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()
        conn.close()
        return dict(row) if row else None


# Singleton
_hosting: Optional[HostingDB] = None


def get_hosting() -> HostingDB:
    global _hosting
    if _hosting is None:
        _hosting = HostingDB(cfg.data_dir / "hosting.db")
    return _hosting