142 lines
4.8 KiB
Python
142 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
One-shot: create a dedicated NewAPI sk- token for the downstream
|
|
"蜂群程序" (swarm bee program) and print the base URL + key.
|
|
|
|
Why direct DB INSERT instead of the /api/token POST endpoint:
|
|
- that endpoint requires a logged-in admin session cookie we
|
|
don't have a clean way to mint headless;
|
|
- the token table shape is stable (new-api fork has tracked
|
|
these columns for two years).
|
|
|
|
Safety:
|
|
- Attaches the token to the platform owner user (whose email
|
|
we look up first).
|
|
- Marks UnlimitedQuota=true so this token consumes user-row
|
|
quota (no separate billing surface to manage).
|
|
- HideFromUserUI=false so the operator can revoke it from the
|
|
/keys page if needed.
|
|
- Sets a descriptive Name + Remark per the user's "备注好"
|
|
requirement.
|
|
"""
|
|
import os
|
|
import sys
|
|
import secrets
|
|
import string
|
|
|
|
import paramiko
|
|
|
|
HOST = "20.24.50.121"
|
|
USER = "heicode"
|
|
PASS = os.environ.get("MGR_PASS", "")
|
|
if not PASS:
|
|
print("MGR_PASS not set", file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
DB_PASS = os.environ.get("MGR_DB_PASS", "")
|
|
DB_HOST = "heicode.postgres.database.azure.com"
|
|
DB_USER = "heicode"
|
|
DB_NAME = "heicode"
|
|
|
|
if not DB_PASS:
|
|
print("MGR_DB_PASS not set", file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
# new-api token keys are 48-char alphanum (no padding). Match that
|
|
# exactly so admin tools / dashboards parse it.
|
|
ALPHABET = string.ascii_letters + string.digits
|
|
RAW_KEY = "".join(secrets.choice(ALPHABET) for _ in range(48))
|
|
# Token bearer format on the wire is "sk-<48 chars>".
|
|
|
|
SQL_FIND = (
|
|
"SELECT id, username, email FROM users "
|
|
"WHERE email IN ('zsbgnw@gmail.com','chenchen@xinghanlab.com') "
|
|
" OR username IN ('chenchen','root') "
|
|
" OR role >= 10 "
|
|
"ORDER BY role DESC, id ASC LIMIT 1;"
|
|
)
|
|
|
|
INSERT_SQL_TEMPLATE = (
|
|
"INSERT INTO tokens "
|
|
'(user_id, name, "key", status, created_time, accessed_time, '
|
|
"expired_time, remain_quota, used_quota, unlimited_quota, "
|
|
"model_limits_enabled, model_limits, allow_ips, "
|
|
'"group", cross_group_retry, hide_from_user_ui) '
|
|
"VALUES "
|
|
"({user_id}, '蜂群程序 (Swarm Bot)', '{key}', 1, "
|
|
"EXTRACT(EPOCH FROM NOW())::bigint, EXTRACT(EPOCH FROM NOW())::bigint, "
|
|
"-1, 0, 0, true, false, '', '', '', false, false) "
|
|
"RETURNING id, name, \"key\";"
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
c = paramiko.SSHClient()
|
|
c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
c.connect(HOST, 22, USER, PASS, look_for_keys=False, allow_agent=False, timeout=30)
|
|
|
|
def psql(sql: str) -> str:
|
|
# Pipe SQL via stdin so quoted identifiers like "group" survive
|
|
# without shell-escape hell.
|
|
cmd = (
|
|
f"PGPASSWORD='{DB_PASS}' psql "
|
|
f"'sslmode=require host={DB_HOST} user={DB_USER} dbname={DB_NAME}' "
|
|
f"-At -F'|' -v ON_ERROR_STOP=1"
|
|
)
|
|
stdin, stdout, stderr = c.exec_command(cmd, timeout=60, get_pty=False)
|
|
stdin.write(sql)
|
|
stdin.channel.shutdown_write()
|
|
out = stdout.read().decode("utf-8", "replace").strip()
|
|
err = stderr.read().decode("utf-8", "replace").strip()
|
|
if err:
|
|
print(f"[psql stderr] {err}", file=sys.stderr)
|
|
return out
|
|
|
|
# Step 1: find owner user
|
|
print(f"--- Looking up platform owner user ---", flush=True)
|
|
owner_row = psql(SQL_FIND)
|
|
if not owner_row:
|
|
print("ERROR: no owner user found", file=sys.stderr)
|
|
sys.exit(3)
|
|
fields = owner_row.split("|")
|
|
owner_id = int(fields[0])
|
|
owner_username = fields[1] if len(fields) > 1 else ""
|
|
owner_email = fields[2] if len(fields) > 2 else ""
|
|
print(f"Owner: id={owner_id} username={owner_username} email={owner_email}")
|
|
|
|
# Step 2: insert the token row
|
|
print(f"\n--- Inserting dedicated 蜂群程序 token ---", flush=True)
|
|
insert_sql = INSERT_SQL_TEMPLATE.format(user_id=owner_id, key=RAW_KEY)
|
|
inserted = psql(insert_sql)
|
|
if not inserted:
|
|
print("ERROR: INSERT returned nothing", file=sys.stderr)
|
|
sys.exit(4)
|
|
fields = inserted.split("|")
|
|
token_id = int(fields[0])
|
|
token_name = fields[1]
|
|
token_key = fields[2]
|
|
print(f"Token row inserted: id={token_id} name={token_name}")
|
|
|
|
# Step 3: print result
|
|
base_url = "https://code.xinghanlab.com"
|
|
print("\n" + "=" * 60)
|
|
print("NEWAPI CREDENTIAL — 蜂群程序 (SWARM BOT)")
|
|
print("=" * 60)
|
|
print(f"Base URL (OpenAI compat): {base_url}/v1")
|
|
print(f"Base URL (Anthropic compat): {base_url}/v1/messages")
|
|
print(f"Base URL (raw root): {base_url}")
|
|
print(f"")
|
|
print(f"API Key (Authorization: Bearer ...):")
|
|
print(f" sk-{token_key}")
|
|
print(f"")
|
|
print(f"Token row id: {token_id}")
|
|
print(f"Owner user id: {owner_id} ({owner_email})")
|
|
print(f"Quota model: unlimited (consumes the owner-account quota)")
|
|
print(f"Expires: never")
|
|
print(f"Manage / revoke from: {base_url}/keys")
|
|
print("=" * 60)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|