From 451c24d3297ab53e37dd6e862ef0d10b1f426d47 Mon Sep 17 00:00:00 2001 From: Nixietab <75538775+nixietab@users.noreply.github.com> Date: Tue, 11 Nov 2025 08:40:32 -0300 Subject: [PATCH] Added cli, better scheme --- main.py | 255 ++++++++++++++++++++---------------------------------- server.py | 183 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 275 insertions(+), 163 deletions(-) create mode 100644 server.py diff --git a/main.py b/main.py index 75b581a..666a11d 100644 --- a/main.py +++ b/main.py @@ -1,174 +1,103 @@ +import argparse import os +import sys import json -import hashlib -import requests -from datetime import datetime, timedelta -from flask import Flask, jsonify, Response, request -from urllib.parse import urljoin -import tokenext +from configparser import ConfigParser from healthcheck import ensure_config_exists -# load config -config = ensure_config_exists() -server_cfg = config["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() +def configparser_to_dict(cp: ConfigParser) -> dict: + d = {} + for section in cp.sections(): + d[section] = {} + for key, val in cp.items(section): + d[section][key] = val + return d -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] -SMN_TOKEN_FILE = "token" -CACHE_TTL = timedelta(minutes=60) -AUTH_ENABLED = PASSWORD != "" - -app = Flask(__name__) - -def log(msg: str): - if LOG_FILE: - with open(LOG_FILE, "a") as f: - f.write(f"[{datetime.now().isoformat()}] {msg}\n") - print(msg) - -# cache handling -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) - - # token handling -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.") +def load_config(config_path=None): + if config_path: + if not os.path.exists(config_path): + print(f"[ERROR] Config file not found: {config_path}") + sys.exit(1) + try: + with open(config_path, "r", encoding="utf-8") as f: + data = json.load(f) + return data + except json.JSONDecodeError: + cp = ConfigParser() + cp.read(config_path) + return configparser_to_dict(cp) + except Exception as e: + print(f"[ERROR] Failed to read config file: {e}") + sys.exit(1) 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 - -# upstream the request -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 - -# the proxy stuff -@app.route(f"{BASE_PATH}/") -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 + cfg = ensure_config_exists() + if isinstance(cfg, ConfigParser): + return configparser_to_dict(cfg) + return cfg + + +def parse_args(): + parser = argparse.ArgumentParser( + prog="OpenSMN", + description="Run the SMN API with custom configuration." + ) + + parser.add_argument( + "-c", "--config", + type=str, + help="Path to a custom configuration file (INI or JSON)." + ) + + parser.add_argument( + "-p", "--port", + type=int, + help="Override the server port defined in config." + ) + + parser.add_argument( + "--no-auth", + action="store_true", + help="Disable authentication even if password is defined in config." + ) + + parser.add_argument( + "--show-config", + action="store_true", + help="Print the loaded configuration and exit." + ) + + parser.add_argument( + "--base-path", + type=str, + help="Override the base API path." + ) + + return parser.parse_args() + + +def main(): + args = parse_args() + config = load_config(args.config) + server_cfg = config.get("server", {}) + + # apply cli overrides + if args.port: + server_cfg["port"] = str(args.port) + if args.base_path: + server_cfg["base_path"] = args.base_path + if args.no_auth: + server_cfg["password"] = "" + + if args.show_config: + print(json.dumps(config, indent=2, ensure_ascii=False)) + sys.exit(0) + + # run the server + import server + server.run_from_cli(config) -# === Startup === if __name__ == "__main__": - 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}/'") - app.run(host="0.0.0.0", port=PORT) + main() diff --git a/server.py b/server.py new file mode 100644 index 0000000..bb5f184 --- /dev/null +++ b/server.py @@ -0,0 +1,183 @@ +import os +import json +import hashlib +import requests +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}/") +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 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}/'") + app.run(host="0.0.0.0", port=PORT) + +if __name__ == "__main__": + run_from_cli() \ No newline at end of file