admin.py

import sys

import db


def usage():
    print("""usage:
  python3 admin.py list
  python3 admin.py add <username> <password> [quota_mb] [rate_kbps]
  python3 admin.py quota <username> <extra_mb>
  python3 admin.py rate <username> <kbps>
  python3 admin.py ban <username>
  python3 admin.py unban <username>
  python3 admin.py reset <username>""")


def set_rate(username, kbps):
    conn = db.get_conn()
    conn.execute("UPDATE users SET rate_bps = ? WHERE username = ?",
                 (kbps * 1024, username))
    conn.commit()
    conn.close()


def reset_usage(username):
    conn = db.get_conn()
    conn.execute("UPDATE users SET used_bytes = 0, status = 'active' WHERE username = ?",
                 (username,))
    conn.commit()
    conn.close()


if len(sys.argv) < 2:
    usage()
    sys.exit()

db.init_db()
cmd = sys.argv[1]

if cmd == "list":
    print(f"{'id':<3} {'username':<12} {'used':>9} {'quota':>9} {'rate':>10} {'status':<10}")
    for u in db.list_users():
        rate = u["rate_bps"] if "rate_bps" in u.keys() else 0
        rate_s = f"{rate//1024}KB/s" if rate else "unlimited"
        print(f"{u['id']:<3} {u['username']:<12} {db.human(u['used_bytes']):>9} "
              f"{db.human(u['quota_bytes']):>9} {rate_s:>10} {u['status']:<10}")

elif cmd == "add":
    name, pwd = sys.argv[2], sys.argv[3]
    quota = int(sys.argv[4]) * 1024 ** 2 if len(sys.argv) > 4 else 2 * 1024 ** 3
    if db.add_user(name, pwd, quota_bytes=quota):
        if len(sys.argv) > 5:
            set_rate(name, int(sys.argv[5]))
        print(f"added {name}")
    else:
        print(f"{name} already exists")

elif cmd == "quota":
    db.add_quota(sys.argv[2], int(sys.argv[3]) * 1024 ** 2)
    print(f"added {sys.argv[3]}MB to {sys.argv[2]}")

elif cmd == "rate":
    set_rate(sys.argv[2], int(sys.argv[3]))
    print(f"rate limit for {sys.argv[2]}: {sys.argv[3]}KB/s")

elif cmd == "ban":
    db.set_status(sys.argv[2], "banned")
    print(f"{sys.argv[2]} banned")

elif cmd == "unban":
    db.set_status(sys.argv[2], "active")
    print(f"{sys.argv[2]} unbanned")

elif cmd == "reset":
    reset_usage(sys.argv[2])
    print(f"{sys.argv[2]} usage reset")

else:
    usage()