forked from FOSS/Steam-Panel
v 3.0.3
This commit is contained in:
Manchik
@@ -0,0 +1,120 @@
|
||||
"""AutoAcceptManager — manages per-account background loops for auto-accepting Steam logins."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import aiohttp
|
||||
from loguru import logger
|
||||
|
||||
from app.services.steam_auto_accept import mobile_login, get_pending_sessions, confirm_session
|
||||
from app.database import get_db
|
||||
|
||||
CHECK_INTERVAL = 15 # seconds
|
||||
|
||||
|
||||
class AutoAcceptManager:
|
||||
def __init__(self) -> None:
|
||||
self._tasks: dict[int, asyncio.Task] = {}
|
||||
self._tokens: dict[int, str] = {}
|
||||
self._errors: dict[int, str] = {}
|
||||
|
||||
async def start(self, account: dict) -> None:
|
||||
account_id = account["id"]
|
||||
if account_id in self._tasks and not self._tasks[account_id].done():
|
||||
return
|
||||
|
||||
self._tasks[account_id] = asyncio.create_task(self._loop(account))
|
||||
logger.info(f"[auto-accept] Started monitoring for {account['login']} (id={account_id})")
|
||||
|
||||
def stop(self, account_id: int) -> None:
|
||||
task = self._tasks.pop(account_id, None)
|
||||
self._tokens.pop(account_id, None)
|
||||
self._errors.pop(account_id, None)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
logger.info(f"[auto-accept] Stopped monitoring for id={account_id}")
|
||||
|
||||
def is_running(self, account_id: int) -> bool:
|
||||
task = self._tasks.get(account_id)
|
||||
return task is not None and not task.done()
|
||||
|
||||
def running_ids(self) -> set[int]:
|
||||
# Cleanup dead tasks
|
||||
dead = [aid for aid, t in self._tasks.items() if t.done()]
|
||||
for aid in dead:
|
||||
self._tasks.pop(aid, None)
|
||||
self._tokens.pop(aid, None)
|
||||
return {aid for aid, t in self._tasks.items() if not t.done()}
|
||||
|
||||
def get_errors(self) -> dict[int, str]:
|
||||
"""Return error messages for accounts that stopped unexpectedly."""
|
||||
return dict(self._errors)
|
||||
|
||||
def pop_errors(self) -> dict[int, str]:
|
||||
"""Return and clear error messages."""
|
||||
errors = dict(self._errors)
|
||||
self._errors.clear()
|
||||
return errors
|
||||
|
||||
async def _loop(self, account: dict) -> None:
|
||||
account_id = account["id"]
|
||||
login = account["login"]
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Initial mobile login
|
||||
token = await mobile_login(session, account)
|
||||
if not token:
|
||||
logger.error(f"[auto-accept] Initial login failed for {login}, stopping")
|
||||
self._errors[account_id] = "Initial login failed"
|
||||
await self._disable_in_db(account_id)
|
||||
return
|
||||
self._tokens[account_id] = token
|
||||
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(CHECK_INTERVAL)
|
||||
|
||||
client_ids = await get_pending_sessions(session, self._tokens[account_id])
|
||||
|
||||
# Token expired — re-login
|
||||
if client_ids is None:
|
||||
logger.warning(f"[auto-accept] Token expired for {login}, re-logging...")
|
||||
new_token = await mobile_login(session, account)
|
||||
if new_token:
|
||||
self._tokens[account_id] = new_token
|
||||
client_ids = await get_pending_sessions(session, new_token)
|
||||
else:
|
||||
logger.error(f"[auto-accept] Re-login failed for {login}, stopping")
|
||||
self._errors[account_id] = "Re-login failed (token expired)"
|
||||
await self._disable_in_db(account_id)
|
||||
return
|
||||
|
||||
if not client_ids:
|
||||
continue
|
||||
|
||||
for cid in client_ids:
|
||||
await confirm_session(session, self._tokens[account_id], account, cid)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"[auto-accept] Loop cancelled for {login}")
|
||||
return
|
||||
except Exception as exc:
|
||||
logger.error(f"[auto-accept] Loop error for {login}: {exc}")
|
||||
await asyncio.sleep(CHECK_INTERVAL)
|
||||
|
||||
async def stop_all(self) -> None:
|
||||
for aid in list(self._tasks):
|
||||
self.stop(aid)
|
||||
|
||||
async def _disable_in_db(self, account_id: int) -> None:
|
||||
"""Reset auto_accept flag in DB when task fails."""
|
||||
try:
|
||||
db = await get_db()
|
||||
await db.execute("UPDATE accounts SET auto_accept = 0 WHERE id = ?", (account_id,))
|
||||
await db.commit()
|
||||
logger.info(f"[auto-accept] Disabled auto_accept in DB for id={account_id}")
|
||||
except Exception as exc:
|
||||
logger.error(f"[auto-accept] Failed to update DB for id={account_id}: {exc}")
|
||||
|
||||
|
||||
auto_accept_manager = AutoAcceptManager()
|
||||
@@ -0,0 +1,20 @@
|
||||
import base64
|
||||
|
||||
import rsa
|
||||
|
||||
|
||||
def get_rsa_key(response_data: dict) -> tuple[str, str, int]:
|
||||
"""Extract RSA key components from Steam's getrsakey response."""
|
||||
mod = response_data["publickey_mod"]
|
||||
exp = response_data["publickey_exp"]
|
||||
timestamp = int(response_data["timestamp"])
|
||||
return mod, exp, timestamp
|
||||
|
||||
|
||||
def encrypt_password(password: str, mod_hex: str, exp_hex: str) -> str:
|
||||
"""RSA-encrypt a password using Steam's public key, return base64."""
|
||||
mod = int(mod_hex, 16)
|
||||
exp = int(exp_hex, 16)
|
||||
public_key = rsa.PublicKey(mod, exp)
|
||||
encrypted = rsa.encrypt(password.encode("ascii"), public_key)
|
||||
return base64.b64encode(encrypted).decode("ascii")
|
||||
@@ -0,0 +1,35 @@
|
||||
class SteamPanelError(Exception):
|
||||
def __init__(self, message: str, code: str = "unknown_error") -> None:
|
||||
self.message = message
|
||||
self.code = code
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class SteamAuthError(SteamPanelError):
|
||||
def __init__(self, message: str = "Authentication failed") -> None:
|
||||
super().__init__(message, code="auth_error")
|
||||
|
||||
|
||||
class SteamGuardError(SteamPanelError):
|
||||
def __init__(self, message: str = "Steam Guard operation failed") -> None:
|
||||
super().__init__(message, code="guard_error")
|
||||
|
||||
|
||||
class ProxyError(SteamPanelError):
|
||||
def __init__(self, message: str = "Proxy error") -> None:
|
||||
super().__init__(message, code="proxy_error")
|
||||
|
||||
|
||||
class MafileError(SteamPanelError):
|
||||
def __init__(self, message: str = "Invalid mafile") -> None:
|
||||
super().__init__(message, code="mafile_error")
|
||||
|
||||
|
||||
class TaskError(SteamPanelError):
|
||||
def __init__(self, message: str = "Task execution failed") -> None:
|
||||
super().__init__(message, code="task_error")
|
||||
|
||||
|
||||
class EmailError(SteamPanelError):
|
||||
def __init__(self, message: str = "Email operation failed") -> None:
|
||||
super().__init__(message, code="email_error")
|
||||
@@ -0,0 +1,138 @@
|
||||
import logging
|
||||
import re
|
||||
import sys
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from app.config import settings
|
||||
|
||||
|
||||
LEVEL_ICONS = {
|
||||
"TRACE": "[~]",
|
||||
"DEBUG": "[.]",
|
||||
"INFO": "[i]",
|
||||
"SUCCESS": "[+]",
|
||||
"WARNING": "[!]",
|
||||
"ERROR": "[e]",
|
||||
"CRITICAL":"[e]",
|
||||
}
|
||||
|
||||
LEVEL_COLORS = {
|
||||
"TRACE": "\033[35m",
|
||||
"DEBUG": "\033[37m",
|
||||
"INFO": "\033[34m",
|
||||
"SUCCESS": "\033[32m",
|
||||
"WARNING": "\033[33m",
|
||||
"ERROR": "\033[31m",
|
||||
"CRITICAL":"\033[31m",
|
||||
}
|
||||
RESET = "\033[0m"
|
||||
|
||||
|
||||
def _loguru_fmt(record: dict) -> str:
|
||||
level = record["level"].name
|
||||
icon = LEVEL_ICONS.get(level, "[i]")
|
||||
color = LEVEL_COLORS.get(level, "")
|
||||
ts = record["time"].strftime("%H:%M:%S")
|
||||
msg = record["message"].replace("<", "\\<").replace("{", "{{").replace("}", "}}")
|
||||
return f"{RESET}[{ts}] {color}{icon}{RESET} {msg}\n"
|
||||
|
||||
|
||||
class _UvicornInterceptHandler(logging.Handler):
|
||||
"""Route all standard logging records (uvicorn, fastapi) into loguru."""
|
||||
|
||||
def emit(self, record: logging.LogRecord) -> None:
|
||||
try:
|
||||
level = logger.level(record.levelname).name
|
||||
except ValueError:
|
||||
level = record.levelno
|
||||
|
||||
frame, depth = sys._getframe(6), 6
|
||||
while frame and frame.f_code.co_filename == logging.__file__:
|
||||
frame = frame.f_back
|
||||
depth += 1
|
||||
|
||||
logger.opt(depth=depth, exception=record.exc_info).log(level, record.getMessage())
|
||||
|
||||
|
||||
_logging_configured = False
|
||||
|
||||
|
||||
def setup_logging() -> None:
|
||||
global _logging_configured
|
||||
if _logging_configured:
|
||||
return
|
||||
_logging_configured = True
|
||||
|
||||
logger.remove()
|
||||
|
||||
logger.add(
|
||||
sys.stderr,
|
||||
format=_loguru_fmt,
|
||||
level="DEBUG" if settings.debug else "INFO",
|
||||
colorize=False,
|
||||
)
|
||||
|
||||
settings.logs_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _file_fmt(record: dict) -> str:
|
||||
level = record["level"].name
|
||||
icon = LEVEL_ICONS.get(level, "[i]")
|
||||
ts = record["time"].strftime("%Y-%m-%d %H:%M:%S")
|
||||
msg = record["message"].replace("<", "\\<").replace("{", "{{").replace("}", "}}")
|
||||
return f"[{ts}] {icon} {msg}\n"
|
||||
|
||||
logger.add(
|
||||
settings.logs_dir / "steampanel_{time:YYYY-MM-DD}.log",
|
||||
format=_file_fmt,
|
||||
level="DEBUG",
|
||||
rotation="1 day",
|
||||
retention="7 days",
|
||||
compression="zip",
|
||||
)
|
||||
|
||||
_install_intercept()
|
||||
|
||||
|
||||
class _AioSQLiteFilter(logging.Filter):
|
||||
"""Reformats raw aiosqlite debug messages into readable DB operation lines."""
|
||||
|
||||
_OP_RE = re.compile(r"method (\w+) of sqlite3\.(\w+)")
|
||||
_SQL_RE = re.compile(r',\s*"(.*?)"(?:,\s*(\(.*?\)))?\s*\)$', re.DOTALL)
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
msg = record.getMessage()
|
||||
# Drop the redundant "operation X completed" lines
|
||||
if " completed" in msg:
|
||||
return False
|
||||
if msg.startswith("executing"):
|
||||
op_m = self._OP_RE.search(msg)
|
||||
method = op_m.group(1).upper() if op_m else "OP"
|
||||
sql_m = self._SQL_RE.search(msg)
|
||||
if sql_m:
|
||||
sql = sql_m.group(1).strip()
|
||||
params = sql_m.group(2) or ""
|
||||
record.msg = f"[DB] {sql}" + (f" {params}" if params else "")
|
||||
else:
|
||||
record.msg = f"[DB] {method}"
|
||||
record.args = ()
|
||||
return True
|
||||
|
||||
|
||||
def _install_intercept() -> None:
|
||||
intercept = _UvicornInterceptHandler()
|
||||
for name in ("", "uvicorn", "uvicorn.error", "uvicorn.access", "fastapi"):
|
||||
log = logging.getLogger(name)
|
||||
log.handlers = [intercept]
|
||||
log.setLevel(logging.DEBUG)
|
||||
log.propagate = False
|
||||
|
||||
# Format aiosqlite SQL logs instead of silencing them
|
||||
aiosqlite_log = logging.getLogger("aiosqlite")
|
||||
aiosqlite_log.setLevel(logging.DEBUG)
|
||||
aiosqlite_log.addFilter(_AioSQLiteFilter())
|
||||
|
||||
logging.getLogger("asyncio").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
import random
|
||||
|
||||
import aiohttp
|
||||
from aiohttp_socks import ProxyConnector
|
||||
from loguru import logger
|
||||
|
||||
from app.database import get_db
|
||||
|
||||
|
||||
def build_proxy_url(address: str, protocol: str) -> str:
|
||||
"""Convert stored address to standard proxy URL (proto://login:pass@ip:port).
|
||||
|
||||
The frontend always normalizes addresses to login:pass@ip:port (or plain ip:port),
|
||||
so we just prepend the scheme.
|
||||
"""
|
||||
if "://" in address:
|
||||
return address
|
||||
scheme = protocol if protocol in ("socks5", "socks4") else "http"
|
||||
return f"{scheme}://{address}"
|
||||
|
||||
|
||||
class ProxyManager:
|
||||
def __init__(self) -> None:
|
||||
self._proxies: list[dict] = []
|
||||
self._index: int = 0
|
||||
|
||||
async def load(self) -> None:
|
||||
db = await get_db()
|
||||
cursor = await db.execute(
|
||||
"SELECT * FROM proxies WHERE is_alive = 1 ORDER BY fail_count ASC"
|
||||
)
|
||||
self._proxies = [dict(r) for r in await cursor.fetchall()]
|
||||
logger.info(f"Loaded {len(self._proxies)} alive proxies")
|
||||
|
||||
def get_next(self) -> dict | None:
|
||||
if not self._proxies:
|
||||
return None
|
||||
proxy = self._proxies[self._index % len(self._proxies)]
|
||||
self._index += 1
|
||||
return proxy
|
||||
|
||||
def get_random(self) -> dict | None:
|
||||
if not self._proxies:
|
||||
return None
|
||||
return random.choice(self._proxies)
|
||||
|
||||
def get_connector(self, proxy: dict | None = None) -> aiohttp.TCPConnector | ProxyConnector:
|
||||
if proxy is None:
|
||||
return aiohttp.TCPConnector()
|
||||
url = build_proxy_url(proxy["address"], proxy.get("protocol", "http"))
|
||||
return ProxyConnector.from_url(url, ssl=False)
|
||||
|
||||
async def mark_failed(self, proxy_id: int) -> None:
|
||||
db = await get_db()
|
||||
await db.execute(
|
||||
"UPDATE proxies SET fail_count = fail_count + 1, last_checked = datetime('now') WHERE id = ?",
|
||||
(proxy_id,),
|
||||
)
|
||||
await db.execute(
|
||||
"UPDATE proxies SET is_alive = 0 WHERE id = ? AND fail_count >= 5",
|
||||
(proxy_id,),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def mark_alive(self, proxy_id: int) -> None:
|
||||
db = await get_db()
|
||||
await db.execute(
|
||||
"UPDATE proxies SET fail_count = 0, is_alive = 1, last_checked = datetime('now') WHERE id = ?",
|
||||
(proxy_id,),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
@property
|
||||
def count(self) -> int:
|
||||
return len(self._proxies)
|
||||
|
||||
|
||||
proxy_manager = ProxyManager()
|
||||
@@ -0,0 +1,249 @@
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Any, Protocol
|
||||
|
||||
from loguru import logger
|
||||
|
||||
from app.config import settings
|
||||
|
||||
DEFAULT_MAX_PARALLEL = 10
|
||||
DEFAULT_ACCOUNT_TIMEOUT = 90 # fallback for actions not in the map below
|
||||
|
||||
# Per-action timeouts (seconds). Fast read-only actions get 60s;
|
||||
# write actions that involve SMS/email prompts get 120s.
|
||||
TASK_TIMEOUTS: dict[str, int] = {
|
||||
"validate": 60,
|
||||
"remove_guard": 60,
|
||||
"change_password": 60,
|
||||
"random_password": 60,
|
||||
"change_phone": 120,
|
||||
"change_email": 120,
|
||||
}
|
||||
|
||||
|
||||
def _read_max_threads() -> int:
|
||||
from app.config import read_validation_settings
|
||||
return int(read_validation_settings().get("max_threads", DEFAULT_MAX_PARALLEL))
|
||||
|
||||
|
||||
def _read_account_timeout() -> int:
|
||||
from app.config import read_validation_settings
|
||||
return int(read_validation_settings().get("account_timeout", DEFAULT_ACCOUNT_TIMEOUT))
|
||||
|
||||
|
||||
class TaskHandler(Protocol):
|
||||
async def __call__(self, account: dict, params: dict, *, task_id: str) -> None: ...
|
||||
|
||||
|
||||
class TaskManager:
|
||||
def __init__(self) -> None:
|
||||
self._running: dict[str, asyncio.Task] = {}
|
||||
self._handlers: dict[str, TaskHandler] = {}
|
||||
self._prompts: dict[str, dict[str, Any]] = {}
|
||||
self._steps: dict[str, dict[str, Any]] = {}
|
||||
self._active_counts: dict[str, int] = {}
|
||||
self._account_results: dict[str, dict[int, dict[str, str]]] = {}
|
||||
self._account_steps: dict[str, dict[int, dict[str, int]]] = {}
|
||||
|
||||
def register_handler(self, action: str, handler: TaskHandler) -> None:
|
||||
self._handlers[action] = handler
|
||||
|
||||
async def submit(
|
||||
self,
|
||||
task_type: str,
|
||||
accounts: list[dict],
|
||||
params: dict,
|
||||
) -> str:
|
||||
task_id = uuid.uuid4().hex[:12]
|
||||
account_ids = [acc.get("id", 0) for acc in accounts]
|
||||
|
||||
from app.database import get_db
|
||||
|
||||
db = await get_db()
|
||||
await db.execute(
|
||||
"INSERT INTO tasks (id, type, status, total, account_ids) VALUES (?, ?, 'running', ?, ?)",
|
||||
(task_id, task_type, len(accounts), json.dumps(account_ids)),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async_task = asyncio.create_task(
|
||||
self._execute(task_id, task_type, accounts, params)
|
||||
)
|
||||
self._running[task_id] = async_task
|
||||
return task_id
|
||||
|
||||
async def _execute(
|
||||
self,
|
||||
task_id: str,
|
||||
task_type: str,
|
||||
accounts: list[dict],
|
||||
params: dict,
|
||||
) -> None:
|
||||
from app.database import get_db
|
||||
|
||||
handler = self._handlers.get(task_type)
|
||||
if not handler:
|
||||
db = await get_db()
|
||||
await db.execute(
|
||||
"UPDATE tasks SET status = 'failed', error = ? WHERE id = ?",
|
||||
(f"No handler for action: {task_type}", task_id),
|
||||
)
|
||||
await db.commit()
|
||||
return
|
||||
|
||||
completed = 0
|
||||
errors: list[str] = []
|
||||
lock = asyncio.Lock()
|
||||
max_parallel = _read_max_threads()
|
||||
account_timeout = TASK_TIMEOUTS.get(task_type, _read_account_timeout())
|
||||
sem = asyncio.Semaphore(max_parallel)
|
||||
self._active_counts[task_id] = 0
|
||||
self._account_results[task_id] = {}
|
||||
total = len(accounts)
|
||||
|
||||
async def process_one(account: dict, idx: int) -> None:
|
||||
nonlocal completed
|
||||
login = account.get("login", "?")
|
||||
acc_id = account.get("id", 0)
|
||||
logger.debug(f"[{task_type}] ({idx}/{total}) processing {login}...")
|
||||
|
||||
async with sem:
|
||||
async with lock:
|
||||
self._active_counts[task_id] = self._active_counts.get(task_id, 0) + 1
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
handler(account, params, task_id=task_id),
|
||||
timeout=account_timeout,
|
||||
)
|
||||
logger.success(f"[{task_type}] ({idx}/{total}) {login} — done")
|
||||
async with lock:
|
||||
self._account_results[task_id][acc_id] = {"status": "ok"}
|
||||
except asyncio.TimeoutError:
|
||||
async with lock:
|
||||
errors.append(f"{login}: timed out after {account_timeout}s")
|
||||
self._account_results[task_id][acc_id] = {"status": "error", "error": f"Timed out after {account_timeout}s"}
|
||||
logger.error(f"[{task_type}] ({idx}/{total}) {login} — timed out after {account_timeout}s")
|
||||
except Exception as exc:
|
||||
err_str = str(exc)
|
||||
async with lock:
|
||||
errors.append(f"{login}: {exc}")
|
||||
self._account_results[task_id][acc_id] = {"status": "error", "error": err_str}
|
||||
logger.error(f"[{task_type}] ({idx}/{total}) {login} — error: {exc}")
|
||||
finally:
|
||||
async with lock:
|
||||
self._active_counts[task_id] = max(0, self._active_counts.get(task_id, 0) - 1)
|
||||
|
||||
async with lock:
|
||||
completed += 1
|
||||
db = await get_db()
|
||||
await db.execute(
|
||||
"UPDATE tasks SET progress = ?, updated_at = datetime('now') WHERE id = ?",
|
||||
(completed, task_id),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
try:
|
||||
tasks = [process_one(acc, i) for i, acc in enumerate(accounts, 1)]
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
db = await get_db()
|
||||
status = "completed"
|
||||
result = f"Done: {completed - len(errors)} success, {len(errors)} errors"
|
||||
error_text = "; ".join(errors[:10]) if errors else None
|
||||
results_json = json.dumps(
|
||||
{str(k): v for k, v in self._account_results.get(task_id, {}).items()}
|
||||
)
|
||||
await db.execute(
|
||||
"UPDATE tasks SET status = ?, result = ?, error = ?, account_results = ?, updated_at = datetime('now') WHERE id = ?",
|
||||
(status, result, error_text, results_json, task_id),
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
except asyncio.CancelledError:
|
||||
db = await get_db()
|
||||
results_json = json.dumps(
|
||||
{str(k): v for k, v in self._account_results.get(task_id, {}).items()}
|
||||
)
|
||||
await db.execute(
|
||||
"UPDATE tasks SET status = 'cancelled', account_results = ?, updated_at = datetime('now') WHERE id = ?",
|
||||
(results_json, task_id),
|
||||
)
|
||||
await db.commit()
|
||||
except Exception as exc:
|
||||
logger.error(f"Task {task_id} fatal error: {exc}")
|
||||
db = await get_db()
|
||||
results_json = json.dumps(
|
||||
{str(k): v for k, v in self._account_results.get(task_id, {}).items()}
|
||||
)
|
||||
await db.execute(
|
||||
"UPDATE tasks SET status = 'failed', error = ?, account_results = ?, updated_at = datetime('now') WHERE id = ?",
|
||||
(str(exc)[:500], results_json, task_id),
|
||||
)
|
||||
await db.commit()
|
||||
finally:
|
||||
self._running.pop(task_id, None)
|
||||
self._steps.pop(task_id, None)
|
||||
self._active_counts.pop(task_id, None)
|
||||
self._account_results.pop(task_id, None)
|
||||
self._account_steps.pop(task_id, None)
|
||||
|
||||
def get_active_count(self, task_id: str) -> int:
|
||||
return self._active_counts.get(task_id, 0)
|
||||
|
||||
def get_account_results(self, task_id: str) -> dict[int, dict[str, str]]:
|
||||
return self._account_results.get(task_id, {})
|
||||
|
||||
def clear_account_results(self, task_id: str) -> None:
|
||||
self._account_results.pop(task_id, None)
|
||||
|
||||
async def set_step(self, task_id: str, step: int, total_steps: int, label: str = "", acc_id: int = 0) -> None:
|
||||
self._steps[task_id] = {"step": step, "total_steps": total_steps, "label": label}
|
||||
if acc_id:
|
||||
self._account_steps.setdefault(task_id, {})[acc_id] = {"step": step, "total": total_steps}
|
||||
|
||||
def get_step_info(self, task_id: str) -> dict[str, Any] | None:
|
||||
return self._steps.get(task_id)
|
||||
|
||||
def get_account_steps(self, task_id: str) -> dict[int, dict[str, int]]:
|
||||
return self._account_steps.get(task_id, {})
|
||||
|
||||
async def prompt_user(self, task_id: str, message: str, login: str = "") -> str:
|
||||
"""Pause task execution and ask the user for input via SSE prompt."""
|
||||
key = f"{task_id}:{login}" if login else task_id
|
||||
event = asyncio.Event()
|
||||
self._prompts[key] = {"message": message, "login": login, "event": event, "response": ""}
|
||||
await event.wait()
|
||||
response = self._prompts.pop(key, {}).get("response", "")
|
||||
return response
|
||||
|
||||
def respond(self, task_id: str, value: str, login: str = "") -> bool:
|
||||
key = f"{task_id}:{login}" if login else task_id
|
||||
prompt = self._prompts.get(key)
|
||||
if not prompt:
|
||||
return False
|
||||
prompt["response"] = value
|
||||
prompt["event"].set()
|
||||
return True
|
||||
|
||||
def get_pending_prompt(self, task_id: str) -> dict[str, str] | None:
|
||||
"""Return first pending prompt for this task (with login info)."""
|
||||
prefix = f"{task_id}:"
|
||||
for key, prompt in self._prompts.items():
|
||||
if key == task_id or key.startswith(prefix):
|
||||
return {"message": prompt["message"], "login": prompt.get("login", "")}
|
||||
return None
|
||||
|
||||
def cancel(self, task_id: str) -> bool:
|
||||
task = self._running.get(task_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
return True
|
||||
return False
|
||||
|
||||
def active_count(self) -> int:
|
||||
return len(self._running)
|
||||
|
||||
|
||||
task_manager = TaskManager()
|
||||
Reference in New Issue
Block a user