access_manager.py

import os
import json
import logging
import random
import string
import psutil
from datetime import datetime, timedelta, timezone
import requests
import functools

def to_aware(dt: datetime) -> datetime:
    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc)
    return dt
import redis
from constants import KEYS_FILE, MASTER_KEY, OUTPUT_DIR

@functools.lru_cache(maxsize=1024)
def get_geo_location(ip: str) -> str:
    if ip in ("127.0.0.1", "::1", "localhost", "Unknown"):
        return ""
    try:
        res = requests.get(f"http://ip-api.com/json/{ip}?fields=status,country,city", timeout=2)
        if res.ok:
            data = res.json()
            if data.get("status") == "success":
                city = data.get("city", "")
                country = data.get("country", "")
                if city and country:
                    return f" ({city}, {country})"
                elif country:
                    return f" ({country})"
    except Exception:
        pass
    return ""

import logging.handlers

# Set up access logger
access_logger = logging.getLogger("wealth_access")
access_logger.setLevel(logging.INFO)
# Rotate every 30 days, keep 1 backup
handler = logging.handlers.TimedRotatingFileHandler(os.path.join(OUTPUT_DIR, "access.log"), when='D', interval=30, backupCount=1, encoding='utf-8')
formatter = logging.Formatter('%(asctime)s - %(message)s')
handler.setFormatter(formatter)
access_logger.addHandler(handler)

# Initialize Redis client if REDIS_URL is provided
redis_url = os.getenv("REDIS_URL")
redis_client = None
if redis_url:
    try:
        redis_client = redis.from_url(redis_url, decode_responses=True)
        access_logger.info("Successfully connected to Redis for access key storage.")
    except Exception as e:
        access_logger.error(f"Failed to connect to Redis: {e}")

def load_keys():
    if redis_client:
        try:
            data = redis_client.get("access_keys")
            if data:
                return json.loads(data)
        except Exception as e:
            access_logger.error(f"Redis read error: {e}")
            
    # Fallback to local JSON
    if not os.path.exists(KEYS_FILE):
        return {"otk": {}}
    try:
        with open(KEYS_FILE, 'r') as f:
            return json.load(f)
    except:
        return {"otk": {}}

def save_keys(data):
    if redis_client:
        try:
            redis_client.set("access_keys", json.dumps(data))
            return
        except Exception as e:
            access_logger.error(f"Redis write error: {e}")
            
    # Fallback to local JSON
    with open(KEYS_FILE, 'w') as f:
        json.dump(data, f, indent=4)

def is_master_key(key: str) -> bool:
    if not key:
        return False
    return key.strip() == MASTER_KEY

def validate_key(key: str, username: str = None, ip: str = "Unknown", silent: bool = False) -> bool:
    if not key:
        return False
        
    key = key.strip()
    if username:
        username = username.strip()
    
    data = load_keys()
    
    if "rate_limits" not in data:
        data["rate_limits"] = {}
        
    now = datetime.now(timezone.utc)
    
    if ip != "Unknown" and ip != "":
        ip_state = data["rate_limits"].get(ip, {"failed": 0, "cooldown_until": None, "cooldown_count": 0, "flagged": False})
        
        if ip_state.get("flagged", False):
            raise ValueError("IP permanently flagged due to excessive failed attempts.")
            
        cooldown_until = ip_state.get("cooldown_until")
        if cooldown_until:
            cooldown_dt = to_aware(datetime.fromisoformat(cooldown_until))
            if now < cooldown_dt:
                raise ValueError("Too many failed attempts. Please try again in 30 minutes.")

    def handle_failure():
        if ip != "Unknown" and ip != "":
            ip_state["failed"] = ip_state.get("failed", 0) + 1
            if ip_state["failed"] >= 5:
                ip_state["failed"] = 0
                ip_state["cooldown_count"] = ip_state.get("cooldown_count", 0) + 1
                ip_state["cooldown_until"] = (now + timedelta(minutes=30)).isoformat()
                
                if ip_state["cooldown_count"] >= 3:
                    ip_state["flagged"] = True
                    try:
                        requests.post(
                            "https://formsubmit.co/ajax/mkoustis05@gmail.com",
                            json={
                                "email": "security@wealthengine.local",
                                "_subject": f"SECURITY ALERT: IP Flagged ({ip})",
                                "message": f"IP Address {ip} has failed to authenticate 15 times and has been permanently flagged."
                            },
                            timeout=5
                        )
                    except Exception as e:
                        access_logger.error(f"Failed to send security alert: {e}")
                        
            data["rate_limits"][ip] = ip_state
            save_keys(data)
            
        if not silent:
            import threading
            def _log_fail():
                geo = get_geo_location(ip)
                msg = f"⛔ **DENIED:** Invalid, expired, or revoked key '{key}' attempted by IP: `{ip}`{geo}"
                access_logger.warning(msg)
            threading.Thread(target=_log_fail, daemon=True).start()
        return False

    def handle_success(template: str):
        if ip != "Unknown" and ip != "":
            if ip in data.get("rate_limits", {}):
                del data["rate_limits"][ip]
                save_keys(data)
        if not silent and template:
            import threading
            def _log_succ():
                geo = get_geo_location(ip)
                msg = template.format(ip=ip, geo=geo, key=key)
                access_logger.info(msg)
            threading.Thread(target=_log_succ, daemon=True).start()
        return True

    # Check Master Key
    if key == MASTER_KEY:
        return handle_success("✅ **MASTER_KEY** used by IP: `{ip}`{geo}")
        
    # Check Time-Based Keys
    if key in data.get("otk", {}) and data["otk"][key].get("revoked", False) == False:
        key_data = data["otk"][key]
        
        # Validate that the key belongs to the requested username
        if username and key_data.get("username") != username:
            return handle_failure()
            
        created_at = to_aware(datetime.fromisoformat(key_data["created_at"]))
        
        if "expires_at" in key_data:
            expires_at = to_aware(datetime.fromisoformat(key_data["expires_at"]))
        else:
            expires_at = created_at + timedelta(hours=1)
            
        if now <= expires_at:
            is_first_use = "used_at" not in key_data
            if is_first_use or key_data.get("used_by_ip") != ip:
                key_data["used_at"] = datetime.now().isoformat()
                key_data["used_by_ip"] = ip
                save_keys(data)
                return handle_success("✅ **Time-Key** '{key}' used by IP: `{ip}`{geo}")
            return handle_success("")
            
    return handle_failure()

def cleanup_expired_keys():
    try:
        data = load_keys()
        if "otk" not in data:
            return
        
        now = datetime.now(timezone.utc)
        cutoff = now - timedelta(days=30)
        
        keys_to_delete = []
        for key, info in data["otk"].items():
            expires_str = info.get("expires_at")
            if expires_str:
                expires_dt = to_aware(datetime.fromisoformat(expires_str))
                if expires_dt < cutoff:
                    keys_to_delete.append(key)
            elif info.get("revoked", False):
                created_str = info.get("created_at")
                if created_str:
                    created_dt = to_aware(datetime.fromisoformat(created_str))
                    if created_dt < cutoff:
                        keys_to_delete.append(key)
                        
        if keys_to_delete:
            for k in keys_to_delete:
                del data["otk"][k]
            save_keys(data)
            access_logger.info(f"Cleaned up {len(keys_to_delete)} old/expired keys.")
    except Exception as e:
        access_logger.error(f"Failed to cleanup keys: {e}")

def generate_otk(admin_key: str, username: str = "anonymous", new_key: str = None, hours: int = 1) -> str:
    if admin_key != MASTER_KEY:
        return None
        
    cleanup_expired_keys()
        
    if not new_key:
        new_key = "OTK-" + ''.join(random.choices(string.ascii_uppercase + string.digits, k=8))
        
    data = load_keys()
    data["otk"][new_key] = {
        "username": username,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "expires_at": (datetime.now(timezone.utc) + timedelta(hours=hours)).isoformat(),
        "revoked": False
    }
    save_keys(data)
    access_logger.info(f"ADMIN: New OTK generated: '{new_key}', valid for {hours} hours.")
    return new_key

def revoke_otk(admin_key: str, target_key: str) -> bool:
    if admin_key != MASTER_KEY:
        return False
        
    data = load_keys()
    if target_key in data["otk"]:
        data["otk"][target_key]["revoked"] = True
        save_keys(data)
        access_logger.info(f"ADMIN: OTK revoked: '{target_key}'")
        return True
    return False

def get_all_keys(admin_key: str) -> dict:
    if admin_key != MASTER_KEY:
        return {}
    cleanup_expired_keys()
    return load_keys().get("otk", {})