db.py

import sqlite3
import hashlib

DB = "users.db"

DEFAULT_QUOTA = 2 * 1024 ** 3


def get_conn():
    conn = sqlite3.connect(DB)
    conn.row_factory = sqlite3.Row
    return conn


def init_db():
    conn = get_conn()
    conn.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            quota_bytes INTEGER DEFAULT 2147483648,
            used_bytes INTEGER DEFAULT 0,
            status TEXT DEFAULT 'active',
            is_admin INTEGER DEFAULT 0,
            created_at TEXT DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.commit()
    conn.close()


def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()


def add_user(username, password, quota_bytes=DEFAULT_QUOTA, is_admin=0):
    conn = get_conn()
    try:
        conn.execute(
            "INSERT INTO users (username, password_hash, quota_bytes, is_admin) VALUES (?, ?, ?, ?)",
            (username, hash_password(password), quota_bytes, is_admin),
        )
        conn.commit()
        return True
    except sqlite3.IntegrityError:
        return False
    finally:
        conn.close()


def get_user(username):
    conn = get_conn()
    row = conn.execute("SELECT * FROM users WHERE username = ?", (username,)).fetchone()
    conn.close()
    return row


def check_login(username, password):
    row = get_user(username)
    if row is None:
        return None
    if row["password_hash"] != hash_password(password):
        return None
    return row


def list_users():
    conn = get_conn()
    rows = conn.execute("SELECT * FROM users ORDER BY id").fetchall()
    conn.close()
    return rows


def add_usage(username, nbytes):
    conn = get_conn()
    conn.execute("UPDATE users SET used_bytes = used_bytes + ? WHERE username = ?",
                 (nbytes, username))
    conn.commit()
    conn.close()


def set_status(username, status):
    conn = get_conn()
    conn.execute("UPDATE users SET status = ? WHERE username = ?", (status, username))
    conn.commit()
    conn.close()


def add_quota(username, extra_bytes):
    conn = get_conn()
    conn.execute(
        "UPDATE users SET quota_bytes = quota_bytes + ?, status = 'active' WHERE username = ?",
        (extra_bytes, username),
    )
    conn.commit()
    conn.close()


def human(n):
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n:.1f}{unit}"
        n /= 1024
    return f"{n:.1f}TB"


if __name__ == "__main__":
    init_db()

    print(add_user("admin", "admin123", is_admin=1))
    print(add_user("daniel", "1234"))
    print(add_user("daniel", "5678"))

    print()
    print(check_login("daniel", "1234") is not None)
    print(check_login("daniel", "wrong") is not None)
    print(check_login("nobody", "1234") is not None)

    print()
    add_usage("daniel", 150 * 1024 * 1024)
    set_status("admin", "banned")

    for u in list_users():
        print(f"{u['id']} {u['username']:<8} {human(u['used_bytes']):>8} / "
              f"{human(u['quota_bytes']):<8} {u['status']:<10} admin={u['is_admin']}")


def log_traffic(username, tun_ip, dst_ip, dst_port, protocol, domain=None):
    conn = get_conn()
    conn.execute(
        "INSERT INTO traffic (username, tun_ip, dst_ip, dst_port, protocol, domain) "
        "VALUES (?, ?, ?, ?, ?, ?)",
        (username, tun_ip, dst_ip, dst_port, protocol, domain),
    )
    conn.commit()
    conn.close()


def get_traffic(username=None, limit=200):
    conn = get_conn()
    if username:
        rows = conn.execute(
            "SELECT * FROM traffic WHERE username = ? ORDER BY id DESC LIMIT ?",
            (username, limit)).fetchall()
    else:
        rows = conn.execute(
            "SELECT * FROM traffic ORDER BY id DESC LIMIT ?", (limit,)).fetchall()
    conn.close()
    return rows


def set_domain(dst_ip, domain):
    conn = get_conn()
    conn.execute("UPDATE traffic SET domain = ? WHERE dst_ip = ? AND domain IS NULL",
                 (domain, dst_ip))
    conn.commit()
    conn.close()


def log_purchase(username, mb):
    conn = get_conn()
    conn.execute("""
        CREATE TABLE IF NOT EXISTS purchases (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT NOT NULL,
            mb INTEGER NOT NULL,
            ts TEXT DEFAULT CURRENT_TIMESTAMP
        )
    """)
    conn.execute("INSERT INTO purchases (username, mb) VALUES (?, ?)", (username, mb))
    conn.commit()
    conn.close()


def get_purchases(username):
    conn = get_conn()
    conn.execute("""
        CREATE TABLE IF NOT EXISTS purchases (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT NOT NULL,
            mb INTEGER NOT NULL,
            ts TEXT DEFAULT CURRENT_TIMESTAMP
        )
    """)
    rows = conn.execute(
        "SELECT * FROM purchases WHERE username = ? ORDER BY id DESC", (username,)
    ).fetchall()
    conn.close()
    return rows