-
Notifications
You must be signed in to change notification settings - Fork 4
Expand file tree
/
Copy pathapi.py
More file actions
1668 lines (1449 loc) · 60.9 KB
/
api.py
File metadata and controls
1668 lines (1449 loc) · 60.9 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
from fastapi import FastAPI, HTTPException, status, UploadFile, Request
from pydantic import BaseModel
from typing import Optional, List, Dict
from sqlalchemy import select, func, delete, desc
from pathlib import Path
from database import get_session, AccountModel, AccountUsageRecordModel, init_db
from fastapi.middleware.cors import CORSMiddleware
from datetime import datetime
import uvicorn
import asyncio
import os
import traceback
from fastapi.responses import JSONResponse, FileResponse, Response
from cursor_pro_keep_alive import main as register_account
from logger import info, error
from contextlib import asynccontextmanager
from tokenManager.cursor import Cursor # 添加这个导入
import concurrent.futures
from functools import lru_cache
from config import (
MAX_ACCOUNTS,
REGISTRATION_INTERVAL,
API_HOST,
API_PORT,
API_DEBUG,
API_WORKERS,
)
from fastapi.staticfiles import StaticFiles
from dotenv import load_dotenv
import json
import time
import uuid
import aiosqlite
import sqlite3
import logging
# 从get_email_code.py导入验证码请求字典
from get_email_code import EmailVerificationHandler
# 修改为直接访问get_email_code模块
import get_email_code
# 全局状态追踪
registration_status = {
"is_running": False,
"last_run": None,
"last_status": None,
"next_run": None,
"total_runs": 0,
"successful_runs": 0,
"failed_runs": 0,
}
# 定义静态文件目录
static_path = Path(__file__).parent / "static"
static_path.mkdir(exist_ok=True) # 确保目录存在
# 全局任务存储
background_tasks = {"registration_task": None}
# 添加用于存储等待验证码的邮箱信息的全局字典
# pending_verification_codes = {}
# 添加lifespan管理器,在应用启动时初始化数据库
@asynccontextmanager
async def lifespan(app: FastAPI):
# 启动时初始化数据库
await init_db()
info(f'服务已启动! \n 首页地址:http://127.0.0.1:{API_PORT}/')
yield
# 关闭时的清理操作
info("应用程序已关闭!")
app = FastAPI(
title="Cursor Account API",
description="API for managing Cursor accounts",
version="1.0.0",
docs_url="/docs",
redoc_url="/redoc",
lifespan=lifespan,
debug=API_DEBUG,
)
# 挂载静态文件目录
app.mount("/static", StaticFiles(directory=str(static_path)), name="static")
# 添加CORS中间件
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
class Account(BaseModel):
email: str
password: Optional[str] = None
token: str
user: str
usage_limit: Optional[str] = None
created_at: Optional[str] = None
status: str = "active" # 默认为"active"
id: Optional[int] = None # 添加id字段,可选
class Config:
from_attributes = True
class AccountResponse(BaseModel):
success: bool
data: Optional[Account] = None
message: str = ""
async def get_active_account_count() -> int:
"""获取当前账号总数"""
async with get_session() as session:
result = await session.execute(
select(func.count())
.select_from(AccountModel)
.where(AccountModel.status == "active")
)
return result.scalar()
async def get_account_count() -> int:
"""获取当前账号总数"""
async with get_session() as session:
result = await session.execute(select(func.count()).select_from(AccountModel))
return result.scalar()
async def run_registration():
"""运行注册脚本"""
global registration_status
browser_manager = None
try:
info("注册任务开始运行")
while registration_status["is_running"]:
try:
count = await get_active_account_count()
info(f"当前数据库已激活账号数: {count}")
if count >= MAX_ACCOUNTS:
# 修改:不再结束任务,而是进入监控模式
info(f"已达到最大账号数量 ({count}/{MAX_ACCOUNTS}),进入监控模式")
registration_status["last_status"] = "monitoring"
# 等待检测间隔时间
next_check = datetime.now().timestamp() + REGISTRATION_INTERVAL
registration_status["next_run"] = next_check
info(f"将在 {REGISTRATION_INTERVAL} 秒后重新检查账号数量")
await asyncio.sleep(REGISTRATION_INTERVAL)
# 跳过当前循环的剩余部分,继续下一次循环检查
continue
info(f"开始注册尝试 (当前账号数: {count}/{MAX_ACCOUNTS})")
registration_status["last_run"] = datetime.now().isoformat()
registration_status["total_runs"] += 1
# 调用注册函数
try:
success = await asyncio.get_event_loop().run_in_executor(
None, register_account
)
if success:
registration_status["successful_runs"] += 1
registration_status["last_status"] = "success"
info("注册成功")
else:
registration_status["failed_runs"] += 1
registration_status["last_status"] = "failed"
info("注册失败")
except SystemExit:
# 捕获 SystemExit 异常,这是注册脚本正常退出的方式
info("注册脚本通过sys.exit退出")
registration_status["successful_runs"] += 1
registration_status["last_status"] = "completed"
except Exception as e:
error(f"注册过程执行出错: {str(e)}")
error(traceback.format_exc())
registration_status["failed_runs"] += 1
registration_status["last_status"] = "error"
# 更新下次运行时间
next_run = datetime.now().timestamp() + REGISTRATION_INTERVAL
registration_status["next_run"] = next_run
info(f"等待 {REGISTRATION_INTERVAL} 秒后进行下一次尝试")
await asyncio.sleep(REGISTRATION_INTERVAL)
except asyncio.CancelledError:
info("注册迭代被取消")
raise
except Exception as e:
registration_status["failed_runs"] += 1
registration_status["last_status"] = "error"
error(f"注册过程出错: {str(e)}")
error(traceback.format_exc())
if not registration_status["is_running"]:
break
await asyncio.sleep(REGISTRATION_INTERVAL)
except asyncio.CancelledError:
info("注册任务被取消")
raise
except Exception as e:
error(f"注册任务致命错误: {str(e)}")
error(traceback.format_exc())
raise
finally:
registration_status["is_running"] = False
if browser_manager:
try:
browser_manager.quit()
except Exception as e:
error(f"清理浏览器资源时出错: {str(e)}")
error(traceback.format_exc())
@app.get("/", tags=["UI"])
async def serve_index():
"""提供Web UI界面"""
index_path = Path(__file__).parent / "index.html"
return FileResponse(index_path)
@app.get("/general", tags=["General"])
async def root():
"""API根路径,返回API信息"""
try:
# 获取当前账号数量和使用情况
async with get_session() as session:
result = await session.execute(select(AccountModel))
accounts = result.scalars().all()
usage_info = []
total_balance = 0
active_accounts = 0
for acc in accounts:
remaining_balance = Cursor.get_remaining_balance(acc.user, acc.token)
remaining_days = Cursor.get_trial_remaining_days(acc.user, acc.token)
if remaining_balance is not None and remaining_balance > 0:
active_accounts += 1
total_balance += remaining_balance
usage_info.append(
{
"email": acc.email,
"balance": remaining_balance,
"days": remaining_days,
"status": (
"active"
if remaining_balance is not None and remaining_balance > 0
else "inactive"
),
}
)
return {
"service": {
"name": "Cursor Account API",
"version": "1.0.0",
"status": "running",
"description": "API for managing Cursor Pro accounts and automatic registration",
},
"statistics": {
"total_accounts": len(accounts),
"active_accounts": active_accounts,
"total_remaining_balance": total_balance,
"max_accounts": MAX_ACCOUNTS,
"remaining_slots": MAX_ACCOUNTS - len(accounts),
"registration_interval": f"{REGISTRATION_INTERVAL} seconds",
},
"accounts_info": usage_info, # 添加账号详细信息
"registration_status": {
"is_running": registration_status["is_running"],
"last_run": registration_status["last_run"],
"last_status": registration_status["last_status"],
"next_run": registration_status["next_run"],
"statistics": {
"total_runs": registration_status["total_runs"],
"successful_runs": registration_status["successful_runs"],
"failed_runs": registration_status["failed_runs"],
"success_rate": (
f"{(registration_status['successful_runs'] / registration_status['total_runs'] * 100):.1f}%"
if registration_status["total_runs"] > 0
else "N/A"
),
},
},
"endpoints": {
"documentation": {"swagger": "/docs", "redoc": "/redoc"},
"health": {
"check": "/health",
"registration_status": "/registration/status",
},
"accounts": {
"list_all": "/accounts",
"random": "/account/random",
"create": {"path": "/account", "method": "POST"},
"delete": {"path": "/account/{email}", "method": "DELETE"},
"usage": {
"path": "/account/{email}/usage",
"method": "GET",
"description": "Get account usage by email",
},
},
"registration": {
"start": {"path": "/registration/start", "method": "GET"},
"stop": {"path": "/registration/stop", "method": "POST"},
"status": {"path": "/registration/status", "method": "GET"},
},
"usage": {"check": {"path": "/usage", "method": "GET"}},
"clean": {
"run": {
"path": "/clean",
"method": "POST",
"params": {"clean_type": ["check", "disable", "delete"]},
}
},
},
"support": {
"github": "https://github.com/Elawen-Carl/cursor-account-api",
"author": "Elawen Carl",
"contact": "elawencarl@gmail.com",
},
"timestamp": datetime.now().isoformat(),
}
except Exception as e:
error(f"根端点错误: {str(e)}")
error(traceback.format_exc())
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="Error fetching API information",
)
@app.get("/health", tags=["General"])
async def health_check():
"""健康检查端点"""
return {"status": "healthy"}
@app.get("/accounts", tags=["Accounts"])
async def get_accounts(
page: int = 1,
per_page: int = 10,
search: Optional[str] = None,
sort_by: str = "created_at", # 新增排序字段参数
order: str = "desc" # 新增排序方向参数
):
"""获取所有账号列表,支持分页、搜索和排序"""
try:
async with get_session() as session:
# 验证排序参数
valid_sort_fields = {"id", "email", "created_at", "usage_limit"}
valid_orders = {"asc", "desc"}
if sort_by not in valid_sort_fields:
sort_by = "created_at" # 默认排序字段
if order not in valid_orders:
order = "desc" # 默认排序方向
# 构建基本查询
query = select(AccountModel)
# 应用排序
sort_field = getattr(AccountModel, sort_by)
if order == "desc":
query = query.order_by(desc(sort_field))
else:
query = query.order_by(sort_field)
# 如果有搜索参数
if search:
search = f"%{search}%"
query = query.where(AccountModel.email.like(search))
# 计算总记录数
count_query = select(func.count()).select_from(AccountModel)
if search:
count_query = count_query.where(AccountModel.email.like(search))
total_count = await session.scalar(count_query)
# 计算分页
total_pages = (total_count + per_page - 1) // per_page # 向上取整
# 应用分页
query = query.offset((page - 1) * per_page).limit(per_page)
# 执行查询
result = await session.execute(query)
accounts = result.scalars().all()
# 转换为可序列化的数据,确保包含id字段
accounts_data = []
for account in accounts:
account_dict = {
"id": account.id, # 确保包含id字段
"email": account.email,
"password": account.password,
"token": account.token,
"user": account.user if hasattr(account, "user") else "",
"usage_limit": account.usage_limit,
"created_at": account.created_at,
"status": account.status
}
accounts_data.append(account_dict)
# 构建分页响应
return {
"success": True,
"data": accounts_data, # 使用序列化后的数据
"pagination": {
"page": page,
"per_page": per_page,
"total_count": total_count,
"total_pages": total_pages
},
"sort": {
"field": sort_by,
"order": order
}
}
except Exception as e:
error(f"获取账号列表失败: {str(e)}")
error(traceback.format_exc())
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"获取账号列表失败: {str(e)}",
)
@app.get("/account/random", response_model=AccountResponse, tags=["Accounts"])
async def get_random_account():
"""随机获取一个可用的账号和token"""
try:
async with get_session() as session:
result = await session.execute(
select(AccountModel).order_by(func.random()).limit(1)
)
account = result.scalar_one_or_none()
if not account:
return AccountResponse(success=False, message="No accounts available")
return AccountResponse(success=True, data=Account.from_orm(account))
except Exception as e:
error(f"获取随机账号失败: {str(e)}")
error(traceback.format_exc())
raise HTTPException(status_code=500, detail="Internal server error")
@app.post("/account", response_model=AccountResponse, tags=["Accounts"])
async def create_account(account: Account):
"""创建新账号"""
try:
async with get_session() as session:
db_account = AccountModel(
email=account.email,
password=account.password,
token=account.token,
usage_limit=account.usage_limit,
created_at=account.created_at,
)
session.add(db_account)
await session.commit()
return AccountResponse(
success=True, data=account, message="Account created successfully"
)
except Exception as e:
error(f"创建账号失败: {str(e)}")
error(traceback.format_exc())
return AccountResponse(
success=False, message=f"Failed to create account: {str(e)}"
)
@app.delete("/account/{email}", response_model=AccountResponse, tags=["Accounts"])
async def delete_account(email: str, hard_delete: bool = False):
"""删除或停用指定邮箱的账号
如果 hard_delete=True,则物理删除账号
否则仅将状态设置为'deleted'
"""
try:
async with get_session() as session:
# 先检查账号是否存在
result = await session.execute(
select(AccountModel).where(AccountModel.email == email)
)
account = result.scalar_one_or_none()
if not account:
return AccountResponse(success=False, message=f"账号 {email} 不存在")
if hard_delete:
# 物理删除账号
await session.execute(
delete(AccountModel).where(AccountModel.email == email)
)
delete_message = f"账号 {email} 已永久删除"
else:
# 逻辑删除:将状态更新为'deleted'
account.status = "deleted"
delete_message = f"账号 {email} 已标记为删除状态"
await session.commit()
return AccountResponse(success=True, message=delete_message)
except Exception as e:
error(f"删除账号失败: {str(e)}")
error(traceback.format_exc())
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"删除账号失败: {str(e)}",
)
# 添加状态更新的请求体模型
class StatusUpdate(BaseModel):
status: str
@app.put("/account/id/{id}/status", response_model=AccountResponse, tags=["Accounts"])
async def update_account_status(id: str, update: StatusUpdate):
"""更新账号状态
可选状态: active (正常), disabled (停用), deleted (已删除)
"""
# 使用update.status替代原先的status参数
try:
# 验证状态值
valid_statuses = ["active", "disabled", "deleted"]
if update.status not in valid_statuses:
return AccountResponse(
success=False,
message=f"无效的状态值。允许的值: {', '.join(valid_statuses)}",
)
async with get_session() as session:
# 通过邮箱查询账号
result = await session.execute(
select(AccountModel).where(AccountModel.id == id)
)
account = result.scalar_one_or_none()
if not account:
return AccountResponse(
success=False, message=f"邮箱为 {email} 的账号不存在"
)
# 更新状态
account.status = update.status
await session.commit()
return AccountResponse(
success=True,
message=f"账号 {account.email} 状态已更新为 '{update.status}'",
)
except Exception as e:
error(f"通过邮箱更新账号状态失败: {str(e)}")
error(traceback.format_exc())
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"更新账号状态失败: {str(e)}",
)
@app.get("/registration/start", tags=["Registration"])
async def start_registration():
"""手动启动注册任务"""
info("手动启动注册任务")
global background_tasks, registration_status
try:
# 检查是否已达到最大账号数
count = await get_active_account_count()
# 检查任务是否已在运行
if (
background_tasks["registration_task"]
and not background_tasks["registration_task"].done()
):
# 确定当前状态
current_status = (
"monitoring"
if registration_status["last_status"] == "monitoring"
else "running"
)
status_message = (
f"注册任务已在运行中 (状态: {current_status})"
if current_status == "running"
else f"已达到最大账号数量({count}/{MAX_ACCOUNTS}),正在监控账号状态,当账号数量减少时将自动继续注册"
)
info(f"注册请求被忽略 - 任务已在{current_status}状态")
return {
"success": True,
"message": status_message,
"status": {
"is_running": registration_status["is_running"],
"last_run": registration_status["last_run"],
"next_run": (
datetime.fromtimestamp(
registration_status["next_run"]
).isoformat()
if registration_status["next_run"]
else None
),
"last_status": registration_status["last_status"],
"current_count": count,
"max_accounts": MAX_ACCOUNTS,
},
}
# 重置注册状态
registration_status.update(
{
"is_running": True,
"last_status": "starting",
"last_run": datetime.now().isoformat(),
"next_run": datetime.now().timestamp() + REGISTRATION_INTERVAL,
"total_runs": 0,
"successful_runs": 0,
"failed_runs": 0,
}
)
# 创建并启动新任务
loop = asyncio.get_running_loop()
task = loop.create_task(run_registration())
background_tasks["registration_task"] = task
# 添加任务完成回调
def task_done_callback(task):
try:
task.result() # 这将重新引发任何未处理的异常
except asyncio.CancelledError:
info("注册任务被取消")
registration_status["last_status"] = "cancelled"
except Exception as e:
error(f"注册任务失败: {str(e)}")
error(traceback.format_exc())
registration_status["last_status"] = "error"
finally:
if registration_status["is_running"]: # 只有在任务仍在运行时才更新状态
registration_status["is_running"] = False
background_tasks["registration_task"] = None
task.add_done_callback(task_done_callback)
info("手动启动注册任务")
# 等待任务实际开始运行
await asyncio.sleep(1)
# 检查任务是否成功启动
if task.done():
try:
task.result() # 如果任务已完成,检查是否有异常
except Exception as e:
error(f"注册任务启动失败: {str(e)}")
error(traceback.format_exc())
registration_status["is_running"] = False
registration_status["last_status"] = "error"
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Failed to start registration task: {str(e)}",
)
# 检查是否已达到最大账号数,预先设置为监控模式
if count >= MAX_ACCOUNTS:
registration_status["last_status"] = "monitoring"
status_message = f"已达到最大账号数量({count}/{MAX_ACCOUNTS}),进入监控模式,当账号数量减少时将自动继续注册"
else:
status_message = "注册任务启动成功"
return {
"success": True,
"message": status_message,
"status": {
"is_running": registration_status["is_running"],
"last_run": registration_status["last_run"],
"next_run": datetime.fromtimestamp(
registration_status["next_run"]
).isoformat(),
"last_status": registration_status["last_status"],
"current_count": count,
"max_accounts": MAX_ACCOUNTS,
},
}
except Exception as e:
error(f"启动注册任务失败: {str(e)}")
error(traceback.format_exc())
registration_status["is_running"] = False
registration_status["last_status"] = "error"
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Failed to start registration task: {str(e)}",
)
@app.get("/registration/stop", tags=["Registration"])
async def stop_registration():
"""手动停止注册任务"""
global background_tasks
try:
if (
not background_tasks["registration_task"]
or background_tasks["registration_task"].done()
):
return {"success": False, "message": "No running registration task found"}
background_tasks["registration_task"].cancel()
try:
await background_tasks["registration_task"]
except asyncio.CancelledError:
info("注册任务被取消")
background_tasks["registration_task"] = None
registration_status["is_running"] = False
registration_status["last_status"] = "manually stopped"
# 清理所有待处理的验证码请求
for email_id, request_info in list(get_email_code.pending_verification_codes.items()):
if request_info["status"] == "pending":
get_email_code.pending_verification_codes.pop(email_id, None)
info(f"已清理待处理的验证码请求: {email_id}")
return {"success": True, "message": "Registration task stopped successfully"}
except Exception as e:
error(f"停止注册任务失败: {str(e)}")
error(traceback.format_exc())
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Failed to stop registration task: {str(e)}",
)
@app.get("/registration/status", tags=["Registration"])
async def get_registration_status():
"""获取注册状态"""
try:
count = await get_account_count()
active_count = await get_active_account_count() # 添加获取活跃账号数
# 更新任务状态逻辑
if (
background_tasks["registration_task"]
and not background_tasks["registration_task"].done()
):
if registration_status["last_status"] == "monitoring":
task_status = "monitoring" # 新增监控状态
else:
task_status = "running"
else:
task_status = "stopped"
status_info = {
"current_count": count,
"active_count": active_count, # 添加活跃账号数
"max_accounts": MAX_ACCOUNTS,
"is_registration_active": active_count < MAX_ACCOUNTS,
"remaining_slots": MAX_ACCOUNTS - active_count,
"task_status": task_status,
"registration_details": {
"is_running": registration_status["is_running"],
"last_run": registration_status["last_run"],
"last_status": registration_status["last_status"],
"next_run": registration_status["next_run"],
"statistics": {
"total_runs": registration_status["total_runs"],
"successful_runs": registration_status["successful_runs"],
"failed_runs": registration_status["failed_runs"],
"success_rate": (
f"{(registration_status['successful_runs'] / registration_status['total_runs'] * 100):.1f}%"
if registration_status["total_runs"] > 0
else "N/A"
),
},
},
}
# 添加状态解释信息
if task_status == "monitoring":
status_info["status_message"] = (
f"已达到最大账号数量({active_count}/{MAX_ACCOUNTS}),正在监控账号状态,当账号数量减少时将自动继续注册"
)
elif task_status == "running":
status_info["status_message"] = (
f"正在执行注册流程,当前账号数:{active_count}/{MAX_ACCOUNTS}"
)
else:
status_info["status_message"] = "注册任务未运行"
# info(f"请求注册状态 (当前账号数: {count}, 活跃账号数: {active_count}, 状态: {task_status})")
return status_info
except Exception as e:
error(f"获取注册状态失败: {str(e)}")
error(traceback.format_exc())
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"Failed to get registration status: {str(e)}",
)
# 自定义异常处理
@app.exception_handler(HTTPException)
async def http_exception_handler(request, exc):
error(f"HTTP错误发生: {exc.detail}")
return JSONResponse(
status_code=exc.status_code, content={"success": False, "message": exc.detail}
)
@app.exception_handler(Exception)
async def general_exception_handler(request, exc):
error(f"意外错误发生: {str(exc)}")
error(f"错误详情: {traceback.format_exc()}")
return JSONResponse(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
content={
"success": False,
"message": "Internal server error occurred",
"detail": str(exc) if app.debug else None,
},
)
# 添加缓存装饰器
@lru_cache(maxsize=100)
def get_account_status(user: str, token: str, timestamp: int):
"""缓存10分钟内的账号状态"""
balance = Cursor.get_remaining_balance(user, token)
days = Cursor.get_trial_remaining_days(user, token)
return {
"balance": balance,
"days": days,
"status": "active" if balance is not None and balance > 0 else "inactive",
}
# 修改 check_usage 接口
@app.get("/usage")
async def check_usage():
try:
async with get_session() as session:
result = await session.execute(select(AccountModel))
accounts = result.scalars().all()
# 使用当前时间的10分钟间隔作为缓存key
cache_timestamp = int(datetime.now().timestamp() / 600)
# 使用线程池并发获取账号状态
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
futures = [
executor.submit(
get_account_status, acc.user, acc.token, cache_timestamp
)
for acc in accounts
]
usage_info = []
for acc, future in zip(accounts, futures):
status = future.result()
usage_info.append(
{
"email": acc.email,
"usage_limit": status["balance"],
"remaining_days": status["days"],
"status": status["status"],
}
)
return {
"total_accounts": len(accounts),
"usage_info": usage_info,
"summary": {
"active_accounts": sum(
1 for info in usage_info if info["status"] == "active"
),
"inactive_accounts": sum(
1 for info in usage_info if info["status"] == "inactive"
),
"total_remaining_balance": sum(
info["usage_limit"] or 0 for info in usage_info
),
},
}
except Exception as e:
error(f"检查使用量失败: {str(e)}")
error(traceback.format_exc())
raise HTTPException(status_code=500, detail=str(e))
@app.get("/account/{email}/usage", tags=["Accounts"])
async def get_account_usage(email: str):
"""根据邮箱查询账户使用量并更新数据库"""
try:
async with get_session() as session:
# 查询指定邮箱的账号
result = await session.execute(
select(AccountModel).where(AccountModel.email == email)
)
account = result.scalar_one_or_none()
if not account:
raise HTTPException(
status_code=404, detail=f"Account with email {email} not found"
)
# 获取账号使用量
remaining_balance = Cursor.get_remaining_balance(
account.user, account.token
)
remaining_days = Cursor.get_trial_remaining_days(
account.user, account.token
)
# 计算总额度和已使用额度
total_limit = 150 # 默认总额度
used_limit = 0
if remaining_balance is not None:
used_limit = total_limit - remaining_balance
if remaining_days is not None and remaining_days == 0:
account.status = "disabled"
# 更新数据库中的usage_limit字段
account.usage_limit = str(remaining_balance)
await session.commit()
db_updated = True
else:
db_updated = False
return {
"success": True,
"email": account.email,
"usage": {
"remaining_balance": remaining_balance,
"total_limit": total_limit,
"used_limit": used_limit,
"remaining_days": remaining_days,
"status": (
"active"
if remaining_balance is not None and remaining_balance > 0
else "inactive"
),
},
"db_updated": db_updated,
"timestamp": datetime.now().isoformat(),
}
except HTTPException:
raise
except Exception as e:
error(f"查询账号使用量失败: {str(e)}")
error(traceback.format_exc())
return {"success": False, "message": f"Failed to get account usage: {str(e)}"}
# 添加通过ID删除账号的API
@app.delete("/account/id/{id}", response_model=AccountResponse, tags=["Accounts"])
async def delete_account_by_id(id: int, hard_delete: bool = False):
"""通过ID删除或停用账号
如果 hard_delete=True,则物理删除账号
否则仅将状态设置为'deleted'
"""
try:
async with get_session() as session:
# 通过ID查询账号
result = await session.execute(
select(AccountModel).where(AccountModel.id == id)
)
account = result.scalar_one_or_none()
if not account:
return AccountResponse(success=False, message=f"ID为 {id} 的账号不存在")
email = account.email # 保存邮箱以在响应中显示
if hard_delete:
# 物理删除账号
await session.execute(delete(AccountModel).where(AccountModel.id == id))
delete_message = f"账号 {email} (ID: {id}) 已永久删除"
else:
# 逻辑删除:将状态更新为'deleted'
account.status = "deleted"
delete_message = f"账号 {email} (ID: {id}) 已标记为删除状态"