socks-manager/models.py

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


# ── 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)
    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)