158 lines
4.1 KiB
Python
158 lines
4.1 KiB
Python
import json
|
|
import os
|
|
import threading
|
|
|
|
try:
|
|
import psycopg
|
|
except ImportError:
|
|
psycopg = None
|
|
|
|
try:
|
|
import psycopg2
|
|
except ImportError:
|
|
psycopg2 = None
|
|
|
|
|
|
_LOCK = threading.Lock()
|
|
|
|
|
|
def _safe_read_json(path):
|
|
if not os.path.exists(path):
|
|
return {}
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def _safe_write_json(path, data):
|
|
tmp_path = f"{path}.tmp"
|
|
with open(tmp_path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|
os.replace(tmp_path, path)
|
|
|
|
|
|
def save_config(path, config):
|
|
with _LOCK:
|
|
_safe_write_json(path, config)
|
|
|
|
|
|
def load_config(path):
|
|
data = _safe_read_json(path)
|
|
if not isinstance(data, dict):
|
|
return {}
|
|
return data
|
|
|
|
|
|
def _to_text(value):
|
|
if value is None:
|
|
return ""
|
|
return str(value)
|
|
|
|
|
|
def _normalize_optional_text(value):
|
|
text = _to_text(value).strip()
|
|
return text or None
|
|
|
|
|
|
def _get_pg_driver():
|
|
if psycopg is not None:
|
|
return "psycopg"
|
|
if psycopg2 is not None:
|
|
return "psycopg2"
|
|
raise RuntimeError("未安装 PostgreSQL 驱动，请先安装 psycopg[binary] 或 psycopg2-binary")
|
|
|
|
|
|
def _get_pg_connection(config):
|
|
driver = _get_pg_driver()
|
|
host = str(config.get("pg_host", "")).strip()
|
|
port = int(config.get("pg_port", 5432) or 5432)
|
|
dbname = str(config.get("pg_db", "")).strip()
|
|
user = str(config.get("pg_user", "")).strip()
|
|
password = str(config.get("pg_password", "")).strip()
|
|
connect_timeout = int(config.get("pg_connect_timeout", 10) or 10)
|
|
|
|
if not host or not dbname or not user:
|
|
raise ValueError("PostgreSQL 配置不完整，请填写 Host、DB、User")
|
|
|
|
kwargs = {
|
|
"host": host,
|
|
"port": port,
|
|
"dbname": dbname,
|
|
"user": user,
|
|
"password": password,
|
|
"connect_timeout": connect_timeout,
|
|
}
|
|
|
|
if driver == "psycopg":
|
|
return psycopg.connect(**kwargs)
|
|
return psycopg2.connect(**kwargs)
|
|
|
|
|
|
def load_accounts(config):
|
|
query = (
|
|
"select email, mail_password, mail_token, chatgpt_password, name, birthdate, created_at "
|
|
"from registered_accounts order by created_at desc nulls last, email asc"
|
|
)
|
|
accounts = []
|
|
|
|
with _get_pg_connection(config) as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(query)
|
|
rows = cur.fetchall()
|
|
|
|
for row in rows:
|
|
email, mail_password, mail_token, chatgpt_password, name, birthdate, created_at = row
|
|
if not email:
|
|
continue
|
|
accounts.append(
|
|
{
|
|
"email": _to_text(email).strip(),
|
|
"mail_password": _to_text(mail_password).strip(),
|
|
"mail_token": _to_text(mail_token).strip(),
|
|
"chatgpt_password": _to_text(chatgpt_password).strip(),
|
|
"name": _to_text(name).strip(),
|
|
"birthdate": _to_text(birthdate).strip(),
|
|
"created_at": _to_text(created_at).strip(),
|
|
"source": "postgres",
|
|
}
|
|
)
|
|
return accounts
|
|
|
|
|
|
def save_account(
|
|
config,
|
|
email,
|
|
mail_password,
|
|
mail_token,
|
|
chatgpt_password=None,
|
|
name=None,
|
|
birthdate=None,
|
|
):
|
|
query = (
|
|
"insert into registered_accounts "
|
|
"(email, mail_password, mail_token, chatgpt_password, name, birthdate) "
|
|
"values (%s, %s, %s, %s, %s, %s) "
|
|
"on conflict (email) do update set "
|
|
"mail_password = excluded.mail_password, "
|
|
"mail_token = excluded.mail_token, "
|
|
"chatgpt_password = excluded.chatgpt_password, "
|
|
"name = excluded.name, "
|
|
"birthdate = excluded.birthdate"
|
|
)
|
|
params = (
|
|
email,
|
|
mail_password,
|
|
mail_token,
|
|
_normalize_optional_text(chatgpt_password),
|
|
_normalize_optional_text(name),
|
|
_normalize_optional_text(birthdate),
|
|
)
|
|
|
|
with _LOCK:
|
|
with _get_pg_connection(config) as conn:
|
|
with conn.cursor() as cur:
|
|
cur.execute(query, params)
|
|
conn.commit()
|