auth.py

import os
import json
import time
import secrets
import threading
from typing import Optional, Dict, Any, List, Tuple

KEYS_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "api_keys.json")
ADMIN_MASTER_KEY = os.environ.get("ADMIN_MASTER_KEY", "secretz_master_admin_2026")
INTERNAL_BOT_KEY = os.environ.get("INTERNAL_BOT_KEY", "sk_internal_telebot_secretz_9921")
DAILY_FREE_LIMIT = int(os.environ.get("DAILY_FREE_LIMIT", "5"))

_lock = threading.Lock()

# Rate limiting sliding windows
# {api_key: [timestamp1, timestamp2, ...]}
_rate_limits: Dict[str, List[float]] = {}

# Public IP Anti-DDoS sliding windows
# {ip: [timestamp1, timestamp2, ...]}
_public_ip_burst: Dict[str, List[float]] = {}

# Public IP Daily usage: {ip: {"date": "YYYY-MM-DD", "count": int}}
_public_ip_daily: Dict[str, Dict[str, Any]] = {}


def load_keys() -> Dict[str, Dict[str, Any]]:
    if os.path.exists(KEYS_FILE):
        try:
            with open(KEYS_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)
                # Ensure all keys have both 'key' and 'api_key' fields
                for k, v in data.items():
                    if "key" not in v:
                        v["key"] = k
                    if "api_key" not in v:
                        v["api_key"] = k
                return data
        except Exception:
            return {}
    return {}


def save_keys(keys_data: Dict[str, Dict[str, Any]]):
    try:
        with open(KEYS_FILE, "w", encoding="utf-8") as f:
            json.dump(keys_data, f, indent=2)
    except Exception as e:
        print(f"Error saving api_keys.json: {e}")


class KeyManager:
    @staticmethod
    def create_key(
        name: str,
        days_valid: Optional[int] = 30,
        rate_limit_per_min: int = 60,
        total_credits: int = 1000
    ) -> Dict[str, Any]:
        """Creates a new API key with expiration and credit limits."""
        api_key = f"sk_live_{secrets.token_hex(16)}"
        now = time.time()
        expires_at = now + (days_valid * 86400) if (days_valid and days_valid > 0) else None

        key_info = {
            "key": api_key,
            "api_key": api_key,
            "name": name,
            "created_at": now,
            "expires_at": expires_at,
            "rate_limit_per_min": rate_limit_per_min,
            "total_credits": total_credits,
            "used_credits": 0,
            "status": "active"
        }

        with _lock:
            keys = load_keys()
            keys[api_key] = key_info
            save_keys(keys)

        return key_info

    @staticmethod
    def list_keys() -> List[Dict[str, Any]]:
        with _lock:
            keys = load_keys()
            result = []
            for k, v in keys.items():
                item = dict(v)
                item["key"] = k
                item["api_key"] = k
                result.append(item)
            return result

    @staticmethod
    def revoke_key(api_key: str) -> bool:
        with _lock:
            keys = load_keys()
            if api_key in keys:
                keys[api_key]["status"] = "revoked"
                save_keys(keys)
                return True
            return False

    @staticmethod
    def delete_key(api_key: str) -> bool:
        with _lock:
            keys = load_keys()
            if api_key in keys:
                del keys[api_key]
                save_keys(keys)
                return True
            return False

    @staticmethod
    def validate_request(api_key: Optional[str], client_ip: str) -> Tuple[bool, int, str, Dict[str, Any]]:
        """
        Validates an incoming request either by API Key or Public Free Tier (IP-based).
        Returns: (is_valid: bool, error_code: int, error_message: str, meta: dict)
        """
        now = time.time()

        # ========================================================
        # Case 1: Master Admin or Internal Bot Key
        # ========================================================
        if api_key in (ADMIN_MASTER_KEY, INTERNAL_BOT_KEY):
            return True, 200, "OK", {"tier": "admin", "remaining": -1}

        # ========================================================
        # Case 2: User provided a custom API Key
        # ========================================================
        if api_key:
            with _lock:
                keys = load_keys()
                if api_key not in keys:
                    return False, 401, "Invalid API Key. Please verify your credentials.", {}

                info = keys[api_key]

                # Status check
                if info.get("status") != "active":
                    return False, 403, "API Key has been revoked or deactivated.", {}

                # Expiry check
                expires_at = info.get("expires_at")
                if expires_at and now > expires_at:
                    info["status"] = "expired"
                    save_keys(keys)
                    return False, 403, "API Key has expired.", {}

                # Credits check
                total_credits = info.get("total_credits", -1)
                used_credits = info.get("used_credits", 0)
                if total_credits != -1 and used_credits >= total_credits:
                    return False, 403, "API Key search credits exhausted.", {}

                # Rate limit (sliding window 60s)
                rate_limit = info.get("rate_limit_per_min", 60)
                if api_key not in _rate_limits:
                    _rate_limits[api_key] = []

                cutoff = now - 60.0
                _rate_limits[api_key] = [t for t in _rate_limits[api_key] if t > cutoff]

                if len(_rate_limits[api_key]) >= rate_limit:
                    retry_after = round(60.0 - (now - _rate_limits[api_key][0]), 1)
                    return False, 429, f"Rate limit exceeded ({rate_limit} req/min). Retry after {retry_after}s.", {}

                # Deduct credit
                _rate_limits[api_key].append(now)
                info["used_credits"] = used_credits + 1
                save_keys(keys)

                remaining = (total_credits - info["used_credits"]) if total_credits != -1 else -1
                return True, 200, "OK", {"tier": "custom", "remaining": remaining}

        # ========================================================
        # Case 3: Public / Anonymous User (Bina Key ke - Daily Free Limit & Anti-DDoS)
        # ========================================================
        with _lock:
            # 1. Anti-DDoS / Rapid Burst Check (Max 1 request per 2 seconds, max 6 per minute)
            if client_ip not in _public_ip_burst:
                _public_ip_burst[client_ip] = []

            cutoff = now - 60.0
            _public_ip_burst[client_ip] = [t for t in _public_ip_burst[client_ip] if t > cutoff]

            if _public_ip_burst[client_ip] and (now - _public_ip_burst[client_ip][-1] < 1.5):
                return False, 429, "Anti-DDoS: Rapid requests detected. Please wait 2 seconds between queries.", {}

            if len(_public_ip_burst[client_ip]) >= 6:
                return False, 429, "Public rate limit exceeded (max 6 req/min). Please slow down.", {}

            # 2. Daily Free Search Limit (Reset at 00:00 UTC)
            today = time.strftime("%Y-%m-%d", time.gmtime())
            if client_ip not in _public_ip_daily or _public_ip_daily[client_ip].get("date") != today:
                _public_ip_daily[client_ip] = {"date": today, "count": 0}

            used_today = _public_ip_daily[client_ip]["count"]
            if used_today >= DAILY_FREE_LIMIT:
                return (
                    False,
                    429,
                    f"Daily free limit reached ({DAILY_FREE_LIMIT}/{DAILY_FREE_LIMIT} used today). Resets at midnight UTC, or enter an Access Key for unlimited searches.",
                    {"tier": "free", "remaining": 0}
                )

            # Record usage
            _public_ip_burst[client_ip].append(now)
            _public_ip_daily[client_ip]["count"] = used_today + 1
            remaining = max(0, DAILY_FREE_LIMIT - _public_ip_daily[client_ip]["count"])

            return True, 200, "OK", {"tier": "free", "remaining": remaining}

    @staticmethod
    def get_public_quota(client_ip: str) -> Dict[str, Any]:
        """Returns current daily quota status for an IP without consuming credit."""
        today = time.strftime("%Y-%m-%d", time.gmtime())
        with _lock:
            record = _public_ip_daily.get(client_ip, {})
            if record.get("date") != today:
                used = 0
            else:
                used = record.get("count", 0)
            return {
                "tier": "free",
                "daily_limit": DAILY_FREE_LIMIT,
                "used_today": used,
                "remaining_today": max(0, DAILY_FREE_LIMIT - used)
            }

    @staticmethod
    def get_key_status(api_key: str) -> Optional[Dict[str, Any]]:
        if api_key == ADMIN_MASTER_KEY:
            return {"name": "Master Admin", "status": "active", "total_credits": -1, "used_credits": 0, "expires_at": None}
        if api_key == INTERNAL_BOT_KEY:
            return {"name": "Telegram Bot Internal", "status": "active", "total_credits": -1, "used_credits": 0, "expires_at": None}

        with _lock:
            keys = load_keys()
            info = keys.get(api_key)
            if info:
                res = dict(info)
                res["key"] = api_key
                res["api_key"] = api_key
                return res
            return None