Added cli, better scheme
This commit is contained in:
Nixietab
parent
9151259f7e
commit
451c24d329
253
main.py
253
main.py
@ -1,174 +1,103 @@
|
|||||||
|
import argparse
|
||||||
import os
|
import os
|
||||||
|
import sys
|
||||||
import json
|
import json
|
||||||
import hashlib
|
from configparser import ConfigParser
|
||||||
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
|
from healthcheck import ensure_config_exists
|
||||||
|
|
||||||
# load config
|
|
||||||
config = ensure_config_exists()
|
|
||||||
server_cfg = config["server"]
|
|
||||||
|
|
||||||
PORT = int(server_cfg.get("port", 6942))
|
def configparser_to_dict(cp: ConfigParser) -> dict:
|
||||||
PASSWORD = server_cfg.get("password", "").strip()
|
d = {}
|
||||||
CACHE_DIR = server_cfg.get("cache_dir", "cache")
|
for section in cp.sections():
|
||||||
BASE_API = server_cfg.get("base_api", "https://ws1.smn.gob.ar")
|
d[section] = {}
|
||||||
LOG_FILE = server_cfg.get("log_file", "").strip()
|
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"
|
def load_config(config_path=None):
|
||||||
CACHE_TTL = timedelta(minutes=60)
|
if config_path:
|
||||||
AUTH_ENABLED = PASSWORD != ""
|
if not os.path.exists(config_path):
|
||||||
|
print(f"[ERROR] Config file not found: {config_path}")
|
||||||
app = Flask(__name__)
|
sys.exit(1)
|
||||||
|
|
||||||
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:
|
try:
|
||||||
with open(path, "r", encoding="utf-8") as f:
|
with open(config_path, "r", encoding="utf-8") as f:
|
||||||
return json.load(f)
|
data = json.load(f)
|
||||||
except Exception:
|
return data
|
||||||
return None
|
except json.JSONDecodeError:
|
||||||
|
cp = ConfigParser()
|
||||||
def save_cache(url: str, data: dict):
|
cp.read(config_path)
|
||||||
os.makedirs(CACHE_DIR, exist_ok=True)
|
return configparser_to_dict(cp)
|
||||||
path = get_cache_filename(url)
|
except Exception as e:
|
||||||
with open(path, "w", encoding="utf-8") as f:
|
print(f"[ERROR] Failed to read config file: {e}")
|
||||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
sys.exit(1)
|
||||||
|
|
||||||
# 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.")
|
|
||||||
else:
|
else:
|
||||||
log("[TOKEN] Failed to refresh token.")
|
cfg = ensure_config_exists()
|
||||||
|
if isinstance(cfg, ConfigParser):
|
||||||
def check_access_token():
|
return configparser_to_dict(cfg)
|
||||||
if not AUTH_ENABLED:
|
return cfg
|
||||||
return True
|
|
||||||
header_token = request.headers.get("Authorization", "").strip()
|
|
||||||
return header_token == PASSWORD
|
def parse_args():
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
# upstream the request
|
prog="OpenSMN",
|
||||||
def fetch_from_smn(url: str, retry: bool = True):
|
description="Run the SMN API with custom configuration."
|
||||||
token = load_smn_token()
|
)
|
||||||
headers = {
|
|
||||||
"Authorization": f"JWT {token}",
|
parser.add_argument(
|
||||||
"Accept": "application/json",
|
"-c", "--config",
|
||||||
"User-Agent": "Mozilla/5.0"
|
type=str,
|
||||||
}
|
help="Path to a custom configuration file (INI or JSON)."
|
||||||
|
)
|
||||||
try:
|
|
||||||
resp = requests.get(url, headers=headers, timeout=10)
|
parser.add_argument(
|
||||||
except requests.RequestException as e:
|
"-p", "--port",
|
||||||
return Response(str(e), status=502)
|
type=int,
|
||||||
|
help="Override the server port defined in config."
|
||||||
if resp.status_code == 401 and retry:
|
)
|
||||||
log("[AUTH] SMN token expired, trying to refresh...")
|
|
||||||
refresh_smn_token()
|
parser.add_argument(
|
||||||
return fetch_from_smn(url, retry=False)
|
"--no-auth",
|
||||||
|
action="store_true",
|
||||||
return resp
|
help="Disable authentication even if password is defined in config."
|
||||||
|
)
|
||||||
# the proxy stuff
|
|
||||||
@app.route(f"{BASE_PATH}/<path:subpath>")
|
parser.add_argument(
|
||||||
def smn_proxy(subpath):
|
"--show-config",
|
||||||
if not check_access_token():
|
action="store_true",
|
||||||
return jsonify({"error": "Unauthorized"}), 401
|
help="Print the loaded configuration and exit."
|
||||||
|
)
|
||||||
if ".." in subpath or subpath.startswith("/"):
|
|
||||||
return jsonify({"error": "Invalid path"}), 400
|
parser.add_argument(
|
||||||
|
"--base-path",
|
||||||
url = urljoin(BASE_API + "/", subpath)
|
type=str,
|
||||||
|
help="Override the base API path."
|
||||||
cached = load_cache(url)
|
)
|
||||||
if cached:
|
|
||||||
log(f"[CACHE] Loaded {subpath}")
|
return parser.parse_args()
|
||||||
return jsonify(cached)
|
|
||||||
|
|
||||||
log(f"[FETCH] {url}")
|
def main():
|
||||||
resp = fetch_from_smn(url)
|
args = parse_args()
|
||||||
|
config = load_config(args.config)
|
||||||
if not hasattr(resp, "status_code"):
|
server_cfg = config.get("server", {})
|
||||||
return Response("Upstream error", status=502)
|
|
||||||
|
# apply cli overrides
|
||||||
if resp.status_code != 200:
|
if args.port:
|
||||||
return Response(resp.text, status=resp.status_code,
|
server_cfg["port"] = str(args.port)
|
||||||
content_type=resp.headers.get("Content-Type", "text/plain"))
|
if args.base_path:
|
||||||
|
server_cfg["base_path"] = args.base_path
|
||||||
try:
|
if args.no_auth:
|
||||||
data = resp.json()
|
server_cfg["password"] = ""
|
||||||
save_cache(url, data)
|
|
||||||
return jsonify(data)
|
if args.show_config:
|
||||||
except Exception:
|
print(json.dumps(config, indent=2, ensure_ascii=False))
|
||||||
return Response("Invalid JSON from SMN", status=502)
|
sys.exit(0)
|
||||||
|
|
||||||
@app.errorhandler(404)
|
# run the server
|
||||||
def handle_not_found(e):
|
import server
|
||||||
return jsonify({
|
server.run_from_cli(config)
|
||||||
"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
|
|
||||||
|
|
||||||
|
|
||||||
# === Startup ===
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
os.makedirs(CACHE_DIR, exist_ok=True)
|
main()
|
||||||
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>'")
|
|
||||||
app.run(host="0.0.0.0", port=PORT)
|
|
||||||
|
|||||||
183
server.py
Normal file
183
server.py
Normal file
@ -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}/<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 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>'")
|
||||||
|
app.run(host="0.0.0.0", port=PORT)
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
run_from_cli()
|
||||||
Loading…
Reference in New Issue
Block a user