socks-manager/models.py

"""SQLAlchemy 数据模型。"""
from datetime import datetime, timezone
from database import db
import bcrypt


# ── 后台管理员 ──────────────────────────────────────────────────
class AdminUser(db.Model):
    """Web 后台管理员账户（存数据库，支持多账户、角色权限）。"""
    __tablename__ = "admin_users"

    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(64), nullable=False, unique=True, index=True)
    password = db.Column(db.String(128), nullable=False)   # bcrypt 哈希
    role = db.Column(db.String(16), default="viewer")       # superadmin | admin | viewer
    enabled = db.Column(db.Boolean, default=True)
    last_login_at = db.Column(db.DateTime, nullable=True)
    last_login_ip = db.Column(db.String(64), nullable=True)
    created_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))
    updated_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc),
                           onupdate=lambda: datetime.now(timezone.utc))

    ROLE_LABELS = {
        "superadmin": "超级管理员",
        "admin": "管理员",
        "viewer": "只读用户",
    }

    # 角色权限位（可扩展）
    ROLE_PERMISSIONS = {
        "superadmin": {"users:read", "users:write", "instances:read", "instances:write",
                       "system:read", "system:write", "admins:read", "admins:write",
                       "logs:read", "stats:read"},
        "admin":      {"users:read", "users:write", "instances:read", "instances:write",
                       "system:read", "system:write", "admins:read",
                       "logs:read", "stats:read"},
        "viewer":     {"users:read", "instances:read", "system:read",
                       "logs:read", "stats:read"},
    }

    def set_password(self, password):
        self.password = bcrypt.hashpw(
            password.encode("utf-8"), bcrypt.gensalt()
        ).decode("utf-8")

    def check_password(self, password):
        if not self.password:
            return False
        try:
            return bcrypt.checkpw(
                password.encode("utf-8"), self.password.encode("utf-8")
            )
        except Exception:
            return False

    @property
    def role_label(self):
        return self.ROLE_LABELS.get(self.role, self.role)

    def has_permission(self, perm):
        perms = self.ROLE_PERMISSIONS.get(self.role, set())
        return perm in perms

    def __repr__(self):
        return f"<AdminUser {self.username} ({self.role})>"


# ── SOCKS5 实例 ──────────────────────────────────────────────────
class Instance(db.Model):
    __tablename__ = "instances"

    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(64), nullable=False, unique=True)
    engine = db.Column(db.String(16), default="builtin")   # builtin / 3proxy
    listen_host = db.Column(db.String(64), default="0.0.0.0")
    listen_port = db.Column(db.Integer, nullable=False)
    # 统一超时(秒):覆盖 SOCKS5 握手/认证/请求读取,以及隧道阶段每段 read。
    # 客户端在任意阶段(含 tunnel)空闲超过此值会被服务端主动断开,避免空连接
    # 占满 fd 直到 LimitNOFILE 上限。
    timeout = db.Column(db.Integer, default=30)
    enabled = db.Column(db.Boolean, default=True)
    notes = db.Column(db.Text)

    # 认证方式
    auth_method = db.Column(db.String(16), default="none")  # none / userpass

    # 流量限速 (Mbps)
    bandwidth_down = db.Column(db.Float, default=0)  # 0=不限
    bandwidth_up = db.Column(db.Float, default=0)

    # 并发限制
    max_concurrent = db.Column(db.Integer, default=10)

    # 创建时间
    created_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))
    updated_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc),
                           onupdate=lambda: datetime.now(timezone.utc))

    # 运行状态 (运行时更新)
    running = db.Column(db.Boolean, default=False)
    active_connections = db.Column(db.Integer, default=0)

    def __repr__(self):
        return f"<Instance {self.name} {self.listen_host}:{self.listen_port}>"


# ── 代理用户 ─────────────────────────────────────────────────────
class User(db.Model):
    __tablename__ = "users"

    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(64), nullable=False, unique=True, index=True)
    password = db.Column(db.String(128), nullable=True)   # 存储 bcrypt 哈希
    enabled = db.Column(db.Boolean, default=True)

    # 流量限制
    total_traffic_mb = db.Column(db.Float, default=0)     # 0=不限
    monthly_traffic_mb = db.Column(db.Float, default=0)   # 0=不限

    # 限速 (Mbps)
    bandwidth_down = db.Column(db.Float, default=0)
    bandwidth_up = db.Column(db.Float, default=0)

    # 并发连接
    max_concurrent = db.Column(db.Integer, default=10)

    # IP 白名单/黑名单（逗号分隔）
    ip_whitelist = db.Column(db.Text, default="")
    ip_blacklist = db.Column(db.Text, default="")

    # 生命周期
    expire_at = db.Column(db.DateTime, nullable=True)  # None=不过期
    banned = db.Column(db.Boolean, default=False)

    # 统计
    bytes_in = db.Column(db.BigInteger, default=0)
    bytes_out = db.Column(db.BigInteger, default=0)

    created_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))
    updated_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc),
                           onupdate=lambda: datetime.now(timezone.utc))

    def set_password(self, password):
        """用 bcrypt 哈希存储密码。"""
        self.password = bcrypt.hashpw(
            password.encode("utf-8"), bcrypt.gensalt()
        ).decode("utf-8")

    def check_password(self, password):
        """验证密码。兼容旧版明文（自动迁移到哈希）。"""
        if not self.password:
            return False
        if self.password.startswith("$2b$") or self.password.startswith("$2a$"):
            try:
                return bcrypt.checkpw(
                    password.encode("utf-8"), self.password.encode("utf-8")
                )
            except Exception:
                return False
        # 兼容旧版明文密码
        if self.password == password:
            self.set_password(password)
            return True
        return False

    def is_active(self):
        """用户是否可用（未过期、未封禁、已启用）

        SQLite 不存 tzinfo, 从数据库读出的 datetime 是 naive;
        我们统一当作 UTC 处理后再比较, 避免 aware/naive TypeError.
        """
        if not self.enabled or self.banned:
            return False
        if self.expire_at:
            exp = self.expire_at
            if exp.tzinfo is None:
                exp = exp.replace(tzinfo=timezone.utc)
            if datetime.now(timezone.utc) > exp:
                return False
        return True

    def __repr__(self):
        return f"<User {self.username}>"


# ── 实时连接 ─────────────────────────────────────────────────────
class Connection(db.Model):
    __tablename__ = "connections"

    id = db.Column(db.String(64), primary_key=True)
    instance_id = db.Column(db.Integer, db.ForeignKey("instances.id"), nullable=True)
    username = db.Column(db.String(64), nullable=True)
    src_ip = db.Column(db.String(64), nullable=True)
    src_port = db.Column(db.Integer, nullable=True)
    dst_addr = db.Column(db.String(256), nullable=True)
    dst_port = db.Column(db.Integer, nullable=True)
    protocol = db.Column(db.String(8), default="SOCKS5")
    state = db.Column(db.String(16), default="active")  # active / closed / rejected
    started_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))
    ended_at = db.Column(db.DateTime, nullable=True)
    bytes_in = db.Column(db.BigInteger, default=0)
    bytes_out = db.Column(db.BigInteger, default=0)

    def __repr__(self):
        return f"<Connection {self.id[:12]} {self.src_ip} -> {self.dst_addr}:{self.dst_port}>"


# ── 审计日志 ─────────────────────────────────────────────────────
class AuditLog(db.Model):
    __tablename__ = "audit_logs"

    id = db.Column(db.Integer, primary_key=True)
    timestamp = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))
    event = db.Column(db.String(32), nullable=False, index=True)
    # event: auth_success, auth_fail, connect_success, connect_reject, user_create,
    #        user_update, instance_start, instance_stop, admin_action
    user = db.Column(db.String(64), nullable=True)
    instance_id = db.Column(db.Integer, nullable=True)
    src_ip = db.Column(db.String(64), nullable=True)
    detail = db.Column(db.Text)

    def __repr__(self):
        return f"<AuditLog {self.event} {self.user}>"


# ── 备份记录 ─────────────────────────────────────────────────────
class Backup(db.Model):
    __tablename__ = "backups"

    id = db.Column(db.Integer, primary_key=True)
    filename = db.Column(db.String(256), nullable=False)
    path = db.Column(db.String(512), nullable=False)
    size = db.Column(db.BigInteger, default=0)
    created_at = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc))
    note = db.Column(db.Text)

    def __repr__(self):
        return f"<Backup {self.filename}>"


# ── 流量快照（每日） ────────────────────────────────────────────
class TrafficSnapshot(db.Model):
    __tablename__ = "traffic_snapshots"

    id = db.Column(db.Integer, primary_key=True)
    date = db.Column(db.Date, nullable=False, unique=True)
    instance_id = db.Column(db.Integer, db.ForeignKey("instances.id"), nullable=True)
    user_id = db.Column(db.Integer, db.ForeignKey("users.id"), nullable=True)
    bytes_in = db.Column(db.BigInteger, default=0)
    bytes_out = db.Column(db.BigInteger, default=0)
    connections = db.Column(db.Integer, default=0)