OpenSMN/server.py

import os
import json
import hashlib
import requests
import signal
import sys
from datetime import datetime, timedelta
from flask import Flask, jsonify, Response, request
from urllib.parse import urljoin
import tokenext
from healthcheck import ensure_config_exists

app = Flask(__name__)

# Global Data
PORT = 6942
PASSWORD = ""
CACHE_DIR = "cache"
BASE_API = "https://ws1.smn.gob.ar"
LOG_FILE = ""
BASE_PATH = "/smn"
AUTH_ENABLED = False
CACHE_TTL = timedelta(minutes=60)
SMN_TOKEN_FILE = "token"

def log(msg: str):
    if LOG_FILE:
        with open(LOG_FILE, "a") as f:
            f.write(f"[{datetime.now().isoformat()}] {msg}\n")
    print(msg)

def get_cache_filename(url: str) -> str:
    h = hashlib.sha256(url.encode()).hexdigest()
    return os.path.join(CACHE_DIR, f"{h}.json")

def load_cache(url: str):
    path = get_cache_filename(url)
    if not os.path.exists(path):
        return None
    mtime = datetime.fromtimestamp(os.path.getmtime(path))
    if datetime.now() - mtime > CACHE_TTL:
        return None
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None

def save_cache(url: str, data: dict):
    os.makedirs(CACHE_DIR, exist_ok=True)
    path = get_cache_filename(url)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

def load_smn_token():
    if not os.path.exists(SMN_TOKEN_FILE):
        log("[TOKEN] Token file not found — refreshing token...")
        refresh_smn_token()
        if not os.path.exists(SMN_TOKEN_FILE):
            raise FileNotFoundError("Token file could not be created.")
    with open(SMN_TOKEN_FILE, "r") as f:
        return f.read().strip()

def refresh_smn_token():
    log("[TOKEN] Refreshing SMN token...")
    ok = tokenext.refresh_token(output_file=SMN_TOKEN_FILE, headless=True, wait_seconds=8)
    if ok:
        log("[TOKEN] Token refreshed successfully.")
    else:
        log("[TOKEN] Failed to refresh token.")

def check_access_token():
    if not AUTH_ENABLED:
        return True
    header_token = request.headers.get("Authorization", "").strip()
    return header_token == PASSWORD

def fetch_from_smn(url: str, retry: bool = True):
    token = load_smn_token()
    headers = {
        "Authorization": f"JWT {token}",
        "Accept": "application/json",
        "User-Agent": "Mozilla/5.0"
    }

    try:
        resp = requests.get(url, headers=headers, timeout=10)
    except requests.RequestException as e:
        return Response(str(e), status=502)

    if resp.status_code == 401 and retry:
        log("[AUTH] SMN token expired, trying to refresh...")
        refresh_smn_token()
        return fetch_from_smn(url, retry=False)

    return resp

@app.route(f"{BASE_PATH}/<path:subpath>")
def smn_proxy(subpath):
    if not check_access_token():
        return jsonify({"error": "Unauthorized"}), 401

    if ".." in subpath or subpath.startswith("/"):
        return jsonify({"error": "Invalid path"}), 400

    url = urljoin(BASE_API + "/", subpath)

    cached = load_cache(url)
    if cached:
        log(f"[CACHE] Loaded {subpath}")
        return jsonify(cached)

    log(f"[FETCH] {url}")
    resp = fetch_from_smn(url)

    if not hasattr(resp, "status_code"):
        return Response("Upstream error", status=502)

    if resp.status_code != 200:
        return Response(resp.text, status=resp.status_code,
                        content_type=resp.headers.get("Content-Type", "text/plain"))

    try:
        data = resp.json()
        save_cache(url, data)
        return jsonify(data)
    except Exception:
        return Response("Invalid JSON from SMN", status=502)

@app.errorhandler(404)
def handle_not_found(e):
    return jsonify({
        "error": "Endpoint not found",
        "message": f"The requested URL '{request.path}' is not a valid API endpoint.",
    }), 200

@app.errorhandler(405)
def handle_method_not_allowed(e):
    return jsonify({
        "error": "Method not allowed",
        "allowed": ["GET"],
        "path": request.path
    }), 200

@app.errorhandler(Exception)
def handle_general_error(e):
    log(f"[ERROR] Unexpected exception: {e}")
    return jsonify({
        "error": "Internal error",
        "message": str(e),
        "path": request.path
    }), 200

def handle_sigint(signum, frame):
    log("shutting down gracefully...")
    sys.exit(0)

signal.signal(signal.SIGINT, handle_sigint)


def run_from_cli(config=None):
    global PORT, PASSWORD, CACHE_DIR, BASE_API, LOG_FILE, BASE_PATH, AUTH_ENABLED

    if config is None:
        config = ensure_config_exists()

    server_cfg = config.get("server", {})
    PORT = int(server_cfg.get("port", 6942))
    PASSWORD = server_cfg.get("password", "").strip()
    CACHE_DIR = server_cfg.get("cache_dir", "cache")
    BASE_API = server_cfg.get("base_api", "https://ws1.smn.gob.ar")
    LOG_FILE = server_cfg.get("log_file", "").strip()
    BASE_PATH = server_cfg.get("base_path", "/smn").strip()

    if not BASE_PATH.startswith("/"):
        BASE_PATH = "/" + BASE_PATH
    if BASE_PATH.endswith("/"):
        BASE_PATH = BASE_PATH[:-1]

    AUTH_ENABLED = PASSWORD != ""

    os.makedirs(CACHE_DIR, exist_ok=True)
    if not os.path.exists(SMN_TOKEN_FILE):
        log("[STARTUP] No token file found, generating a new one.")
        refresh_smn_token()

    log(f"[STARTUP] Server starting on port {PORT}")
    log(f"[STARTUP] Base path set to '{BASE_PATH}/<path>'")

    try:
        app.run(host="0.0.0.0", port=PORT)
    except KeyboardInterrupt:
        log("[SHUTDOWN] Interrupted. server stopped cleanly.")
    except Exception as e:
        log(f"[ERROR] Unhandled exception: {e}")
    finally:
        log("[SHUTDOWN] Goodbye.")

if __name__ == "__main__":
    import signal
    import sys

    def handle_sigint(signum, frame):
        log("[SHUTDOWN] shutting down gracefully...")
        sys.exit(0)

    signal.signal(signal.SIGINT, handle_sigint)
    run_from_cli()