Stock-Tool/license_check.py

"""
License Check
=============
Validates the stored license key against the Verimund API on every launch.
Falls back to a local cache if the server is unreachable (30-minute grace period).

Storage location: %APPDATA%\\StockScreener\\
  license.dat    — XOR-encrypted license key (machine-bound)
  auth_cache.dat — last successful auth result + timestamp + JWT token
"""

import base64
import hashlib
import hmac as _hmac_module
import json
import os
import time

import requests

# ---------------------------------------------------------------------------
# API config
# ---------------------------------------------------------------------------
_API_URL = "https://api.verimundsolutions.com"

# HMAC secret — must match HMAC_SECRET in /srv/api/.env on the server.
# This value is obfuscated by PyArmor before distribution.
_HMAC_SECRET = "fa1fac57c9cc9742ac66542a17b032ccdfb2211b72c81fccd1617f17ce03ac7f"

# 30 minutes
_GRACE_PERIOD_SECONDS = 30 * 60

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
_APP_DATA     = os.path.join(os.environ.get("APPDATA", os.path.expanduser("~")), "StockScreener")
_LICENSE_FILE = os.path.join(_APP_DATA, "license.dat")
_CACHE_FILE   = os.path.join(_APP_DATA, "auth_cache.dat")


def _ensure_dir() -> None:
    os.makedirs(_APP_DATA, exist_ok=True)


# ---------------------------------------------------------------------------
# Encryption helpers (XOR cipher keyed on HWID — machine-binds the stored key)
# ---------------------------------------------------------------------------

def _derive_key(hwid: str) -> bytes:
    return hashlib.sha256(hwid.encode("utf-8")).digest()


def _xor(data: bytes, key: bytes) -> bytes:
    return bytes(b ^ key[i % len(key)] for i, b in enumerate(data))


def _encrypt(text: str, hwid: str) -> str:
    raw = _xor(text.encode("utf-8"), _derive_key(hwid))
    return base64.urlsafe_b64encode(raw).decode("ascii")


def _decrypt(token: str, hwid: str) -> str:
    raw = base64.urlsafe_b64decode(token.encode("ascii"))
    return _xor(raw, _derive_key(hwid)).decode("utf-8")


# ---------------------------------------------------------------------------
# License key persistence
# ---------------------------------------------------------------------------

def save_license_key(key: str, hwid: str) -> None:
    """Encrypt and persist the license key to disk."""
    _ensure_dir()
    token = _encrypt(key, hwid)
    with open(_LICENSE_FILE, "w", encoding="ascii") as f:
        f.write(token)


def load_license_key(hwid: str) -> str | None:
    """Return the stored license key, or None if not found / corrupt."""
    if not os.path.exists(_LICENSE_FILE):
        return None
    try:
        with open(_LICENSE_FILE, "r", encoding="ascii") as f:
            token = f.read().strip()
        return _decrypt(token, hwid)
    except Exception:
        return None


def delete_license_key() -> None:
    """Remove stored license key (e.g. on deactivation)."""
    if os.path.exists(_LICENSE_FILE):
        os.remove(_LICENSE_FILE)
    if os.path.exists(_CACHE_FILE):
        os.remove(_CACHE_FILE)


# ---------------------------------------------------------------------------
# Auth cache
# ---------------------------------------------------------------------------

def _save_cache(valid: bool, reason: str, channel: str = "stable", jwt_token: str = "") -> None:
    _ensure_dir()
    payload = {
        "valid": valid,
        "reason": reason,
        "channel": channel,
        "jwt_token": jwt_token,
        "timestamp": time.time(),
    }
    with open(_CACHE_FILE, "w", encoding="utf-8") as f:
        json.dump(payload, f)


def _load_cache() -> dict | None:
    if not os.path.exists(_CACHE_FILE):
        return None
    try:
        with open(_CACHE_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def get_channel() -> str:
    """Return the update channel ('stable' or 'beta') from the last auth cache."""
    cache = _load_cache()
    if cache:
        return cache.get("channel", "stable")
    return "stable"


def get_jwt() -> str:
    """Return the stored JWT token, or empty string if not available."""
    cache = _load_cache()
    if cache:
        return cache.get("jwt_token", "")
    return ""


def get_api_headers() -> dict:
    """Return headers for authenticated API requests."""
    return {
        "Authorization": f"Bearer {get_jwt()}",
        "Content-Type": "application/json",
    }


# ---------------------------------------------------------------------------
# HMAC request signing
# ---------------------------------------------------------------------------

def _sign_request(license_key: str, timestamp: str) -> str:
    """Generate HMAC-SHA256 signature for an auth request."""
    return _hmac_module.new(
        _HMAC_SECRET.encode(),
        f"{license_key}{timestamp}".encode(),
        hashlib.sha256,
    ).hexdigest()


# ---------------------------------------------------------------------------
# Core verification
# ---------------------------------------------------------------------------

def verify_license(license_key: str, hwid: str) -> tuple[bool, str]:
    """
    Returns (valid: bool, reason: str).

    Calls POST /auth on the Verimund API with a signed request.
    Falls back to cached result within the 30-minute grace period if offline.
    """
    try:
        timestamp = str(int(time.time()))
        signature = _sign_request(license_key, timestamp)

        resp = requests.post(
            f"{_API_URL}/auth",
            json={
                "license_key": license_key,
                "hwid": hwid,
                "timestamp": timestamp,
                "signature": signature,
            },
            timeout=10,
        )
        resp.raise_for_status()
        data = resp.json()
        token = data.get("token", "")
        channel = data.get("channel", "stable")
        _save_cache(True, "ok", channel, token)
        return True, "ok"

    except requests.exceptions.HTTPError as e:
        try:
            detail = e.response.json().get("detail", str(e))
        except Exception:
            detail = str(e)
        _save_cache(False, detail, "stable", "")
        return False, detail

    except requests.exceptions.SSLError as e:
        return False, f"SSL error contacting license server: {e}"

    except requests.exceptions.ConnectionError as e:
        cache = _load_cache()
        if cache and cache.get("valid"):
            age = time.time() - cache.get("timestamp", 0)
            if age < _GRACE_PERIOD_SECONDS:
                return True, "offline_grace"
        return False, f"Cannot connect to license server. Check your internet connection.\n\n({e})"

    except requests.RequestException as e:
        cache = _load_cache()
        if cache and cache.get("valid"):
            age = time.time() - cache.get("timestamp", 0)
            if age < _GRACE_PERIOD_SECONDS:
                return True, "offline_grace"
        return False, f"License check failed: {e}"