pryscraper/account_manager.py

# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Rug Munch Media LLC
#
# Part of Pry — https://git.rugmunch.io/RugMunchMedia/pryscraper
# Licensed under MIT. See LICENSE.
"""Pry — Account pool management, session persistence, proxy scoring."""

import json
import logging
import os
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

ACCOUNTS_DIR = Path(os.path.expanduser("~/.pry/accounts"))
ACCOUNTS_DIR.mkdir(parents=True, exist_ok=True)


class AccountPool:
    """Manage pool of registered accounts with session persistence."""

    def store(self, site: str, credentials: dict[str, Any], profile_id: str = "", metadata: dict | None = None) -> str:
        import uuid
        account_id = uuid.uuid4().hex[:12]
        account = {
            "id": account_id, "site": site, "credentials": credentials,
            "profile_id": profile_id, "metadata": metadata or {},
            "status": "active", "created_at": datetime.now(UTC).isoformat(),
            "last_used": None, "use_count": 0, "errors": [],
        }
        path = ACCOUNTS_DIR / f"{site}_{account_id}.json"
        path.write_text(json.dumps(account, indent=2))
        logger.info("account_stored", extra={"site": site, "account_id": account_id})
        return account_id

    def get_active(self, site: str) -> dict[str, Any] | None:
        """Get a random active account for a site."""
        accounts = []
        for path in ACCOUNTS_DIR.glob(f"{site}_*.json"):
            try:
                acct = json.loads(path.read_text())
                if acct.get("status") == "active":
                    accounts.append(acct)
            except (json.JSONDecodeError, OSError):
                continue
        if not accounts:
            return None
        import random
        return random.choice(accounts)

    def mark_error(self, account_id: str, error: str) -> None:
        for path in ACCOUNTS_DIR.glob(f"*_{account_id}.json"):
            try:
                acct = json.loads(path.read_text())
                acct["errors"].append({"time": datetime.now(UTC).isoformat(), "error": error[:200]})
                if len(acct["errors"]) > 5:
                    acct["status"] = "suspended"
                path.write_text(json.dumps(acct, indent=2))
            except Exception:
                pass

    def record_use(self, account_id: str) -> None:
        for path in ACCOUNTS_DIR.glob(f"*_{account_id}.json"):
            try:
                acct = json.loads(path.read_text())
                acct["last_used"] = datetime.now(UTC).isoformat()
                acct["use_count"] = acct.get("use_count", 0) + 1
                path.write_text(json.dumps(acct, indent=2))
            except Exception:
                pass

    def list_accounts(self, site: str = "") -> list[dict[str, Any]]:
        accounts = []
        pattern = f"{site}_*.json" if site else "*_*.json"
        for path in sorted(ACCOUNTS_DIR.glob(pattern), key=os.path.getmtime, reverse=True)[:50]:
            try:
                data = json.loads(path.read_text())
                data.pop("credentials", None)  # Don't expose passwords
                accounts.append(data)
            except Exception:
                continue
        return accounts


class ProxyScorer:
    """Score and rank proxies for reliability."""

    async def test_proxy(self, proxy_url: str, test_url: str = "https://httpbin.org/ip", timeout: int = 10) -> dict[str, Any]:
        from client import get_client
        client = await get_client()
        start = time.time()
        try:
            resp = await client.get(test_url, timeout=timeout)
            elapsed = time.time() - start
            return {"proxy": proxy_url, "working": resp.is_success, "latency": round(elapsed, 2),
                    "status": resp.status_code, "ip": resp.text[:50] if resp.is_success else ""}
        except Exception as e:
            return {"proxy": proxy_url, "working": False, "error": str(e)[:80]}

    def score(self, test_result: dict[str, Any]) -> int:
        """Score proxy 0-100."""
        if not test_result.get("working"):
            return 0
        latency = test_result.get("latency", 99)
        if latency < 1:
            return 100
        if latency < 2:
            return 90
        if latency < 5:
            return 75
        if latency < 10:
            return 50
        return 25