Add monthly quotas, per-IP limits, user management CLI
Quota system now supports: - Per-user: daily, weekly, monthly limits - Per-IP: daily, weekly, monthly limits (abuse prevention) - Per-IP connection limit - Max test duration New CLI flags: --monthly-quota, --ip-daily, --ip-weekly, --ip-monthly User management subcommands: btest-server-pro useradd <user> <pass> btest-server-pro userdel <user> btest-server-pro userlist btest-server-pro userset <user> --enabled true/false --daily N --weekly N New DB tables: ip_usage (per-IP daily tracking) New methods: get_monthly_usage, get_ip_*_usage, start/end_session, delete_user, set_user_enabled, set_user_quota Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -61,6 +61,16 @@ impl UserDb {
|
||||
UNIQUE(username, date)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS ip_usage (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
ip TEXT NOT NULL,
|
||||
date TEXT NOT NULL,
|
||||
tx_bytes INTEGER DEFAULT 0,
|
||||
rx_bytes INTEGER DEFAULT 0,
|
||||
test_count INTEGER DEFAULT 0,
|
||||
UNIQUE(ip, date)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS sessions (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
username TEXT NOT NULL,
|
||||
@@ -74,6 +84,7 @@ impl UserDb {
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_usage_user_date ON usage(username, date);
|
||||
CREATE INDEX IF NOT EXISTS idx_ip_usage_date ON ip_usage(ip, date);
|
||||
CREATE INDEX IF NOT EXISTS idx_sessions_peer ON sessions(peer_ip, started_at);
|
||||
")?;
|
||||
Ok(())
|
||||
@@ -166,6 +177,127 @@ impl UserDb {
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub fn get_monthly_usage(&self, username: &str) -> anyhow::Result<(u64, u64)> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let result = conn.query_row(
|
||||
"SELECT COALESCE(SUM(tx_bytes),0), COALESCE(SUM(rx_bytes),0) FROM usage
|
||||
WHERE username = ?1 AND date >= date('now', '-30 days')",
|
||||
params![username],
|
||||
|row| {
|
||||
let a: i64 = row.get(0)?;
|
||||
let b: i64 = row.get(1)?;
|
||||
Ok((a as u64, b as u64))
|
||||
},
|
||||
)?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
// --- Per-IP usage tracking ---
|
||||
|
||||
pub fn record_ip_usage(&self, ip: &str, tx_bytes: u64, rx_bytes: u64) -> anyhow::Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let today = chrono_date_today();
|
||||
conn.execute(
|
||||
"INSERT INTO ip_usage (ip, date, tx_bytes, rx_bytes, test_count)
|
||||
VALUES (?1, ?2, ?3, ?4, 1)
|
||||
ON CONFLICT(ip, date) DO UPDATE SET
|
||||
tx_bytes = tx_bytes + ?3,
|
||||
rx_bytes = rx_bytes + ?4,
|
||||
test_count = test_count + 1",
|
||||
params![ip, today, tx_bytes as i64, rx_bytes as i64],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn get_ip_daily_usage(&self, ip: &str) -> anyhow::Result<(u64, u64)> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let today = chrono_date_today();
|
||||
let result = conn.query_row(
|
||||
"SELECT COALESCE(SUM(tx_bytes),0), COALESCE(SUM(rx_bytes),0) FROM ip_usage WHERE ip = ?1 AND date = ?2",
|
||||
params![ip, today],
|
||||
|row| {
|
||||
let a: i64 = row.get(0)?;
|
||||
let b: i64 = row.get(1)?;
|
||||
Ok((a as u64, b as u64))
|
||||
},
|
||||
)?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub fn get_ip_weekly_usage(&self, ip: &str) -> anyhow::Result<(u64, u64)> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let result = conn.query_row(
|
||||
"SELECT COALESCE(SUM(tx_bytes),0), COALESCE(SUM(rx_bytes),0) FROM ip_usage
|
||||
WHERE ip = ?1 AND date >= date('now', '-7 days')",
|
||||
params![ip],
|
||||
|row| {
|
||||
let a: i64 = row.get(0)?;
|
||||
let b: i64 = row.get(1)?;
|
||||
Ok((a as u64, b as u64))
|
||||
},
|
||||
)?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub fn get_ip_monthly_usage(&self, ip: &str) -> anyhow::Result<(u64, u64)> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let result = conn.query_row(
|
||||
"SELECT COALESCE(SUM(tx_bytes),0), COALESCE(SUM(rx_bytes),0) FROM ip_usage
|
||||
WHERE ip = ?1 AND date >= date('now', '-30 days')",
|
||||
params![ip],
|
||||
|row| {
|
||||
let a: i64 = row.get(0)?;
|
||||
let b: i64 = row.get(1)?;
|
||||
Ok((a as u64, b as u64))
|
||||
},
|
||||
)?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
// --- Session tracking ---
|
||||
|
||||
pub fn start_session(&self, username: &str, peer_ip: &str, protocol: &str, direction: &str) -> anyhow::Result<i64> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute(
|
||||
"INSERT INTO sessions (username, peer_ip, protocol, direction) VALUES (?1, ?2, ?3, ?4)",
|
||||
params![username, peer_ip, protocol, direction],
|
||||
)?;
|
||||
Ok(conn.last_insert_rowid())
|
||||
}
|
||||
|
||||
pub fn end_session(&self, session_id: i64, tx_bytes: u64, rx_bytes: u64) -> anyhow::Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute(
|
||||
"UPDATE sessions SET ended_at = datetime('now'), tx_bytes = ?1, rx_bytes = ?2 WHERE id = ?3",
|
||||
params![tx_bytes as i64, rx_bytes as i64, session_id],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn delete_user(&self, username: &str) -> anyhow::Result<bool> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let rows = conn.execute("DELETE FROM users WHERE username = ?1", params![username])?;
|
||||
Ok(rows > 0)
|
||||
}
|
||||
|
||||
pub fn set_user_enabled(&self, username: &str, enabled: bool) -> anyhow::Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute(
|
||||
"UPDATE users SET enabled = ?1 WHERE username = ?2",
|
||||
params![enabled as i32, username],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn set_user_quota(&self, username: &str, daily: i64, weekly: i64, monthly: i64) -> anyhow::Result<()> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
conn.execute(
|
||||
"UPDATE users SET daily_quota = ?1, weekly_quota = ?2 WHERE username = ?3",
|
||||
params![daily, weekly, username],
|
||||
)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn list_users(&self) -> anyhow::Result<Vec<User>> {
|
||||
let conn = self.conn.lock().unwrap();
|
||||
let mut stmt = conn.prepare(
|
||||
|
||||
Reference in New Issue
Block a user