diff --git a/files/nats_registration_listener-0225.py b/files/nats_registration_listener-0225.py new file mode 100644 index 0000000..740aeb2 --- /dev/null +++ b/files/nats_registration_listener-0225.py @@ -0,0 +1,711 @@ +#!/usr/bin/env python3 +""" +NATS Registration Listener (fox100 + NetBox hostname/action_next lookup + timing + problem counter) +------------------------------------------------------------------------------------------- +- One device GET (status + tags + custom_fields.action_next), no duplicate fetch +- Keeps: nb_problems counter, timings, iface_id diagnostics, same formatting +- For fox100: + * If action_next is empty/absent -> print a simple stdout note and do nothing else + * If action_next present and action_last != action_next -> publish task_name=action_next + and on success set action_last=action_next and action_next_timestamp=now_epoch + * If action_next present and action_last == action_next -> publish only if + (now_epoch - action_next_timestamp) >= 600; if timestamp missing/invalid -> allow publish + and on success set action_last=action_next and action_next_timestamp=now_epoch + * If action_next present -> prepend 3x ASCII BEL to stdout line (kept behavior) +""" + +import argparse +import asyncio +import json +import os +import signal +import ssl +import sys +import hashlib +import time +from time import monotonic +from datetime import datetime, timezone +from typing import Optional, Dict, Any, Tuple, Set +from urllib.parse import urlencode +from urllib.request import Request, urlopen +from urllib.error import URLError, HTTPError +import base64 # for RabbitMQ Basic Auth + +import nats + + +# ========================= +# NetBox hardcoded config +# ========================= +NB_URL = "http://netbox.gt-tiso.ikeja.co.za" # Base URL +NB_TOKEN = "7648e4f5ee370cda7834682e61b47c2ee8e95623" # keep as provided +NB_TIMEOUT = 3.0 # seconds per HTTP GET + +# Cache for MAC -> NetBox lookup result (seconds). Keeps NetBox load down under chatty devices. +NB_LOOKUP_CACHE_TTL = float(os.environ.get("NB_LOOKUP_CACHE_TTL", "10.0")) +NB_LOOKUP_CACHE = {} # mac_norm -> (expires_monotonic, cached_tuple) + +# ========================= +# RabbitMQ hardcoded config +# ========================= +RMQ_HOST = "10.210.12.2" +RMQ_PORT = 15672 +RMQ_USER = "admin" +RMQ_PASS = "change_me" +RMQ_VHOST = "app" +RMQ_EXCHANGE_WORK = "deviceconfig" # direct exchange (immediate) +RMQ_EXCHANGE_DELAYED = "deviceconfig.delayed" # delayed exchange (x-delayed-message) +RMQ_ROUTING_KEY = "deviceconfig" +RMQ_TIMEOUT = 3.0 + +# ---- Human-editable delay (milliseconds). Set to 0 to disable delay. +# Example: 600000 = 10 minutes +RMQ_DELAY_MS = 15000 + +# Product -> Tag slug mapping (kept unchanged, though not used now) +PRODUCT_TAG_SLUG = { + "fox100": "fox100-auto-upgrade-latest", + # "fox200": "fox200-auto-upgrade-latest", +} + +# Global counter for any NetBox-related problems +NB_PROBLEM_COUNTER = 0 +NB_PROBLEM_LOCK = asyncio.Lock() + + +# ========================= +# Arg parsing +# ========================= +def parse_args(): + p = argparse.ArgumentParser(description="Listen to a NATS subject and print one line per device.") + p.add_argument("--servers", nargs="+", default=["nats://127.0.0.1:4222"]) + p.add_argument("--subject", default="registration") + p.add_argument("--queue", default=None) + p.add_argument("--name", default="registration-listener") + p.add_argument("--creds") + p.add_argument("--user") + p.add_argument("--password") + p.add_argument("--token") + p.add_argument("--tls-ca") + p.add_argument("--tls-cert") + p.add_argument("--tls-key") + p.add_argument("--insecure", action="store_true") + p.add_argument("--include-subject", action="store_true") + return p.parse_args() + + +# ========================= +# Helpers +# ========================= +def make_ssl_context(args) -> Optional[ssl.SSLContext]: + if not any([args.tls_ca, args.tls_cert, args.tls_key]) and not any(url.startswith("tls://") for url in args.servers): + return None + ctx = ssl.create_default_context(purpose=ssl.Purpose.SERVER_AUTH) + if args.tls_ca: + ctx.load_verify_locations(args.tls_ca) + if args.tls_cert and args.tls_key: + ctx.load_cert_chain(certfile=args.tls_cert, keyfile=args.tls_key) + if args.insecure: + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + return ctx + + +def ts() -> str: + return datetime.now(timezone.utc).astimezone().strftime("%Y-%m-%d %H:%M:%S%z") + + +def pick_first_interface(eths: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if "eth0" in eths and isinstance(eths["eth0"], dict): + return eths["eth0"] + for name in sorted(eths.keys()): + if isinstance(eths[name], dict): + return eths[name] + return None + + +def extract_fields(obj: Dict[str, Any]): + root = obj + d = root["data"] if isinstance(root.get("data"), dict) else root + product = d.get("productName") or "-" + fw_active = d.get("firmwareVersion", {}).get("active") or "-" + mac = "-" + eths = d.get("ethernetInterfaces", {}) + if isinstance(eths, dict): + chosen = pick_first_interface(eths) + if chosen and isinstance(chosen.get("macAddress"), str): + mac = chosen["macAddress"] + return product, mac, fw_active + + +# ========================= +# NetBox lookup (urllib) +# ========================= +def normalize_mac(mac: str) -> Optional[str]: + if not mac or not isinstance(mac, str): + return None + s = mac.strip().lower().replace("-", ":") + hex_only = "".join(ch for ch in s if ch in "0123456789abcdef") + if len(hex_only) == 12: + return ":join".replace(":", "").join([":".join(hex_only[i:i+2] for i in range(0, 12, 2))]) # (keeping original behavior; no change) + parts = s.split(":") + if len(parts) == 6 and all(len(p) == 2 and all(c in "0123456789abcdef" for c in p) for p in parts): + return s + return None + + +def http_get_json(url: str, params: Optional[Dict[str, str]] = None, headers: Optional[Dict[str, str]] = None, timeout: float = NB_TIMEOUT): + if params: + url = f"{url}?{urlencode(params)}" + req = Request(url, headers=headers or {}, method="GET") + try: + with urlopen(req, timeout=timeout) as resp: + if resp.status != 200: + return None, resp.status + data = resp.read() + return json.loads(data.decode("utf-8", errors="replace")), 200 + except HTTPError as e: + return None, getattr(e, "code", 599) + except URLError: + return None, 598 + except Exception: + return None, 597 + + +# RabbitMQ management API POST helper (basic auth; JSON in/out) +def http_post_json(url: str, payload_obj: Dict[str, Any], user: Optional[str] = None, password: Optional[str] = None, timeout: float = RMQ_TIMEOUT): + body = json.dumps(payload_obj).encode("utf-8") + headers = {"Content-Type": "application/json"} + if user and password: + token = base64.b64encode(f"{user}:{password}".encode("utf-8")).decode("ascii") + headers["Authorization"] = f"Basic {token}" + req = Request(url, data=body, headers=headers, method="POST") + try: + with urlopen(req, timeout=timeout) as resp: + data = resp.read() + try: + return json.loads(data.decode("utf-8", errors="replace")), resp.status + except Exception: + return None, resp.status + except HTTPError as e: + return None, getattr(e, "code", 599) + except URLError: + return None, 598 + except Exception: + return None, 597 + + +# NetBox PATCH helper (JSON in/out) +def http_patch_json(url: str, payload_obj: Dict[str, Any], headers: Optional[Dict[str, str]] = None, timeout: float = NB_TIMEOUT): + body = json.dumps(payload_obj).encode("utf-8") + h = dict(headers or {}) + h["Content-Type"] = "application/json" + req = Request(url, data=body, headers=h, method="PATCH") + try: + with urlopen(req, timeout=timeout) as resp: + data = resp.read() + try: + return json.loads(data.decode("utf-8", errors="replace")), resp.status + except Exception: + return None, resp.status + except HTTPError as e: + return None, getattr(e, "code", 599) + except URLError: + return None, 598 + except Exception: + return None, 597 + + +async def nb_problem(log_status, msg: str): + """Increment counter and log a problem line.""" + global NB_PROBLEM_COUNTER + async with NB_PROBLEM_LOCK: + NB_PROBLEM_COUNTER += 1 + await log_status(msg) + + +def nb_lookup_device_by_mac(mac: str, log_status) -> Tuple[ + Optional[str], Optional[int], Optional[int], Optional[str], Optional[Set[str]], Optional[Any], Optional[Any], Optional[Any], Optional[Any] +]: + """ + Resolve MAC -> (hostname, iface_id, device_id, device_status_value, tag_slugs_set, action_next, action_last, action_next_timestamp, action_state) + - Logs problems for anomalies (mac not found, unassigned, wrong type, iface fetch fail). + - If device detail fetch fails, returns host/id with status/tags/custom_fields as None. + """ + mac_norm = normalize_mac(mac) + if not mac_norm: + return None, None, None, None, None, None, None, None, None + + # Step 0: short TTL cache (avoid repeated NetBox GETs for chatty devices) + if NB_LOOKUP_CACHE_TTL > 0: + nowm = monotonic() + cached = NB_LOOKUP_CACHE.get(mac_norm) + if cached: + exp, val = cached + if exp > nowm: + host, iface_id, dev_id, status_val, tag_slugs, action_next, action_last, action_next_timestamp, action_state = val + if tag_slugs is not None: + try: + tag_slugs = set(tag_slugs) + except Exception: + pass + return host, iface_id, dev_id, status_val, tag_slugs, action_next, action_last, action_next_timestamp, action_state + NB_LOOKUP_CACHE.pop(mac_norm, None) + + base = NB_URL.rstrip("/") + h = { + "Accept": "application/json", + "Content-Type": "application/json", + "Authorization": f"Token {NB_TOKEN}", + } + + # Step 1: MAC lookup + data, code = http_get_json( + f"{base}/api/dcim/mac-addresses/", + params={"mac_address": mac_norm, "limit": "2", "fields": "assigned_object_type,assigned_object_id"}, + headers=h, + ) + if code == 400: + data, code = http_get_json(f"{base}/api/dcim/mac-addresses/", params={"mac_address": mac_norm, "limit": "2"}, headers=h) + if code != 200 or not data: + asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: mac query http={code} mac={mac_norm}")) + return None, None, None, None, None, None, None, None, None + + results = (data or {}).get("results") or [] + if not results: + asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: mac not found mac={mac_norm}")) + return None, None, None, None, None, None, None, None, None + + rec = results[0] + aot = (rec.get("assigned_object_type") or "").strip() + aoid = rec.get("assigned_object_id") + + if len(results) > 1: + asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: multiple mac records mac={mac_norm} iface_id={aoid if aoid is not None else '-'}")) + + if not aot or aoid is None: + asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: mac unassigned mac={mac_norm}")) + return None, None, None, None, None, None, None, None, None + if aot != "dcim.interface": + asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: mac assigned to {aot} mac={mac_norm} iface_id={aoid}")) + return None, aoid, None, None, None, None, None, None, None + + # Step 2: Interface -> Device (shallow) + iface, code2 = http_get_json(f"{base}/api/dcim/interfaces/{aoid}/", params={"fields": "device"}, headers=h) + if code2 == 400: + iface, code2 = http_get_json(f"{base}/api/dcim/interfaces/{aoid}/", headers=h) + if code2 != 200 or not iface: + asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: iface fetch http={code2} iface_id={aoid}")) + return None, aoid, None, None, None, None, None, None, None + + dev = iface.get("device") or {} + host = dev.get("name") or dev.get("display") + dev_id = dev.get("id") + if not host or dev_id is None: + asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: iface has no device iface_id={aoid}")) + return None, aoid, None, None, None, None, None, None, None + + # Step 3: Device detail (single fetch for status, tags, custom_fields.*) + device, code3 = http_get_json(f"{base}/api/dcim/devices/{dev_id}/", params={"fields": "status,tags,custom_fields"}, headers=h) + if code3 == 400: + device, code3 = http_get_json(f"{base}/api/dcim/devices/{dev_id}/", headers=h) + if code3 != 200 or not device: + # treat as "no extra info" + return host, aoid, dev_id, None, None, None, None, None, None + + status_val = ((device.get("status") or {}).get("value")) or None + tags = device.get("tags") or [] + tag_slugs = set() + for t in tags: + slug = t.get("slug") + if isinstance(slug, str): + tag_slugs.add(slug) + + cf = device.get("custom_fields") or {} + action_next = cf.get("action_next") + action_last = cf.get("action_last") + action_next_timestamp = cf.get("action_next_timestamp") + action_state = cf.get("action_state") + + # Step 4: populate cache (only on full success) + if NB_LOOKUP_CACHE_TTL > 0: + try: + NB_LOOKUP_CACHE[mac_norm] = ( + monotonic() + NB_LOOKUP_CACHE_TTL, + ( + host, + aoid, + dev_id, + (status_val if isinstance(status_val, str) else None), + (tuple(tag_slugs) if tag_slugs is not None else None), + action_next, + action_last, + action_next_timestamp, + action_state, + ), + ) + except Exception: + pass + + return host, aoid, dev_id, (status_val if isinstance(status_val, str) else None), tag_slugs, action_next, action_last, action_next_timestamp, action_state + + +# ========================= +# Main +# ========================= +async def main(): + args = parse_args() + ssl_ctx = make_ssl_context(args) + + print_lock = asyncio.Lock() + + async def log_status(s: str): + async with print_lock: + print(s, file=sys.stderr, flush=True) + + DEDUPE_TTL = float(os.environ.get("DEDUP_TTL", "2.0")) + recent_payloads: Dict[bytes, float] = {} + + async def disconnected_cb(): + await log_status(f"[{ts()}] Disconnected from NATS.") + + async def reconnected_cb(): + await log_status(f"[{ts()}] Reconnected to NATS.") + + async def error_cb(e): + await log_status(f"[{ts()}] Error: {e!r}") + + async def closed_cb(): + await log_status(f"[{ts()}] Connection closed.") + + nc = await nats.connect( + servers=args.servers, + name=args.name, + allow_reconnect=True, + reconnect_time_wait=2, + max_reconnect_attempts=-1, + disconnected_cb=disconnected_cb, + reconnected_cb=reconnected_cb, + error_cb=error_cb, + closed_cb=closed_cb, + user_credentials=args.creds if args.creds else None, + user=args.user, + password=args.password, + token=args.token, + tls=ssl_ctx, + ) + + async def message_handler(msg: nats.aio.msg.Msg): + t_start = monotonic() + payload = msg.data + + # strict payload dedupe + digest = hashlib.blake2b(payload, digest_size=16).digest() + nowm = monotonic() + exp = recent_payloads.get(digest) + if exp and exp > nowm: + return + recent_payloads[digest] = nowm + DEDUPE_TTL + if len(recent_payloads) > 4096: + cutoff = nowm + for k in list(recent_payloads.keys()): + if recent_payloads[k] <= cutoff: + recent_payloads.pop(k, None) + + product = mac = fw = "-" + try: + text = payload.decode("utf-8", errors="replace") + obj = json.loads(text) + product, mac, fw = extract_fields(obj) + except Exception: + pass + + host_suffix = "" + action_suffix = "" # kept; not used + bell_prefix = "" # ASCII BEL when action_next present (3x) + netbox_time_ms = 0.0 + + if product == "fox100": + nb_start = time.perf_counter() + try: + host, iface_id, dev_id, status_val, tag_slugs, action_next, action_last, action_next_timestamp, action_state = nb_lookup_device_by_mac( + mac=mac, log_status=log_status + ) + if host: + host_suffix = f" host={host}" + if iface_id and not host: + host_suffix += f" iface_id={iface_id}" + + if host: + # Gate on action_state: allow only "" or "ready" + action_state_str = "" + try: + if action_state is None: + action_state_str = "" + elif isinstance(action_state, str): + action_state_str = action_state.strip() + else: + action_state_str = str(action_state).strip() + except Exception: + action_state_str = "" + + if action_state_str not in ("", "ready", "done"): + async with print_lock: + print( + f"[{ts()}] device is not ready because of action_state host={host} action_state={action_state_str}", + file=sys.stdout, + flush=True, + ) + return + + # Determine if action_next is present (non-empty string, or any truthy value) + has_action_next = False + action_next_str = None + try: + if isinstance(action_next, str): + action_next_str = action_next.strip() + has_action_next = len(action_next_str) > 0 + else: + has_action_next = bool(action_next) + if has_action_next: + action_next_str = str(action_next) + except Exception: + has_action_next = False + action_next_str = None + + if not has_action_next: + # User-requested behavior: if no action -> just shoot a message to stdout and we're ok + async with print_lock: + print(f"[{ts()}] no action_next for host={host}", file=sys.stdout, flush=True) + else: + # Compare action_last with action_next (strings) + action_last_str = None + try: + if isinstance(action_last, str): + action_last_str = action_last.strip() + elif action_last is None: + action_last_str = None + else: + action_last_str = str(action_last) + except Exception: + action_last_str = None + + now_epoch = int(time.time()) + + # If action_last == action_next, apply cooldown based on action_next_timestamp (600s) + if action_last_str == action_next_str: + allow_repeat = True + try: + if action_next_timestamp is None: + allow_repeat = True + elif isinstance(action_next_timestamp, (int, float)): + allow_repeat = (now_epoch - int(action_next_timestamp)) >= 600 + elif isinstance(action_next_timestamp, str): + allow_repeat = (now_epoch - int(action_next_timestamp.strip())) >= 600 + else: + allow_repeat = True + except Exception: + allow_repeat = True + + if not allow_repeat: + async with print_lock: + print( + f"[{ts()}] cooldown action_next for host={host} task={action_next_str}", + file=sys.stdout, + flush=True, + ) + else: + # Publish task_name=action_next + effective_delay_ms = RMQ_DELAY_MS + target_exchange = RMQ_EXCHANGE_DELAYED if effective_delay_ms > 0 else RMQ_EXCHANGE_WORK + rmq_url = f"http://{RMQ_HOST}:{RMQ_PORT}/api/exchanges/{RMQ_VHOST}/{target_exchange}/publish" + + payload_obj = { + "inscope_device": host, + "task_name": action_next_str, + } + + payload_raw = json.dumps(payload_obj, separators=(",", ":"), ensure_ascii=False) + + rmq_body = { + "properties": { + "content_type": "application/json" + }, + "routing_key": RMQ_ROUTING_KEY, + "payload": payload_raw, + "payload_encoding": "string", + } + + if effective_delay_ms > 0: + rmq_body["properties"]["headers"] = {"x-delay": int(effective_delay_ms)} + + await log_status( + f"[{ts()}] ok, here i will execute\n" + f" url: {rmq_url}\n" + f" routing_key: {RMQ_ROUTING_KEY}\n" + f" payload_raw: {payload_raw}\n" + f" publish_body: {json.dumps(rmq_body, ensure_ascii=False)}" + ) + + resp, code = http_post_json(rmq_url, rmq_body, user=RMQ_USER, password=RMQ_PASS, timeout=RMQ_TIMEOUT) + if code != 200: + await log_status(f"[{ts()}] rmq: publish http={code} host={host}") + else: + routed = False + try: + routed = bool((resp or {}).get("routed", False)) + except Exception: + routed = False + + # ---- SURGICAL FIX ---- + # For delayed publishes (effective_delay_ms > 0), routed may be false but the message is accepted. + publish_ok = True + if effective_delay_ms <= 0 and not routed: + publish_ok = False + await log_status(f"[{ts()}] rmq: publish immediate routed=false host={host}") + # ---------------------- + + if publish_ok: + # On success: set action_last and action_next_timestamp and action_state + try: + base = NB_URL.rstrip("/") + nb_headers = { + "Accept": "application/json", + "Authorization": f"Token {NB_TOKEN}", + } + patch_body = {"custom_fields": {"action_last": action_next_str, "action_next_timestamp": str(now_epoch), "action_state": "started"}} + _, pcode = http_patch_json( + f"{base}/api/dcim/devices/{dev_id}/", + patch_body, + headers=nb_headers, + timeout=NB_TIMEOUT, + ) + if pcode != 200: + await log_status(f"[{ts()}] nb: action_last/timestamp patch http={pcode} host={host} dev_id={dev_id}") + except Exception as e: + await log_status(f"[{ts()}] nb: action_last/timestamp patch error host={host!r} dev_id={dev_id!r} err={e!r}") + + bell_prefix = "\a" * 3 + else: + # action_last != action_next -> publish + effective_delay_ms = RMQ_DELAY_MS + target_exchange = RMQ_EXCHANGE_DELAYED if effective_delay_ms > 0 else RMQ_EXCHANGE_WORK + rmq_url = f"http://{RMQ_HOST}:{RMQ_PORT}/api/exchanges/{RMQ_VHOST}/{target_exchange}/publish" + + payload_obj = { + "inscope_device": host, + "task_name": action_next_str, + } + + payload_raw = json.dumps(payload_obj, separators=(",", ":"), ensure_ascii=False) + + rmq_body = { + "properties": { + "content_type": "application/json" + }, + "routing_key": RMQ_ROUTING_KEY, + "payload": payload_raw, + "payload_encoding": "string", + } + + if effective_delay_ms > 0: + rmq_body["properties"]["headers"] = {"x-delay": int(effective_delay_ms)} + + await log_status( + f"[{ts()}] ok, here i will execute\n" + f" url: {rmq_url}\n" + f" routing_key: {RMQ_ROUTING_KEY}\n" + f" payload_raw: {payload_raw}\n" + f" publish_body: {json.dumps(rmq_body, ensure_ascii=False)}" + ) + + resp, code = http_post_json(rmq_url, rmq_body, user=RMQ_USER, password=RMQ_PASS, timeout=RMQ_TIMEOUT) + if code != 200: + await log_status(f"[{ts()}] rmq: publish http={code} host={host}") + else: + routed = False + try: + routed = bool((resp or {}).get("routed", False)) + except Exception: + routed = False + + # ---- SURGICAL FIX ---- + publish_ok = True + if effective_delay_ms <= 0 and not routed: + publish_ok = False + await log_status(f"[{ts()}] rmq: publish immediate routed=false host={host}") + # ---------------------- + + if publish_ok: + now_epoch = int(time.time()) + try: + base = NB_URL.rstrip("/") + nb_headers = { + "Accept": "application/json", + "Authorization": f"Token {NB_TOKEN}", + } + patch_body = {"custom_fields": {"action_last": action_next_str, "action_next_timestamp": str(now_epoch), "action_state": "started"}} + _, pcode = http_patch_json( + f"{base}/api/dcim/devices/{dev_id}/", + patch_body, + headers=nb_headers, + timeout=NB_TIMEOUT, + ) + if pcode != 200: + await log_status(f"[{ts()}] nb: action_last/timestamp patch http={pcode} host={host} dev_id={dev_id}") + except Exception as e: + await log_status(f"[{ts()}] nb: action_last/timestamp patch error host={host!r} dev_id={dev_id!r} err={e!r}") + + bell_prefix = "\a" * 3 + + except Exception as e: + await nb_problem(log_status, f"[{ts()}] nb: unexpected error mac={mac!r} err={e!r}") + netbox_time_ms = (time.perf_counter() - nb_start) * 1000 + + total_ms = (monotonic() - t_start) * 1000 + + async with NB_PROBLEM_LOCK: + nb_problems_snapshot = NB_PROBLEM_COUNTER + + line = f"nb_problems={nb_problems_snapshot} [{ts()}] product={product} mac={mac} fw={fw}{host_suffix}{action_suffix}" + if product == "fox100": + line += f" netbox_ms={netbox_time_ms:.1f} total_ms={total_ms:.1f}" + if args.include_subject: + line += f" subject={msg.subject}" + + async with print_lock: + sys.stdout.write(bell_prefix + line + "\n") + sys.stdout.flush() + + if args.queue: + await nc.subscribe(args.subject, queue=args.queue, cb=message_handler) + else: + await nc.subscribe(args.subject, cb=message_handler) + + await log_status( + f"[{ts()}] Listening on subject '{args.subject}' (queue={args.queue or '-'}) via {args.servers} | nb=on" + ) + + stop_event = asyncio.Event() + + def handle_signal(*_): + asyncio.create_task(log_status(f"[{ts()}] Received stop signal, draining...")) + stop_event.set() + + loop = asyncio.get_running_loop() + for s in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(s, handle_signal) + except NotImplementedError: + signal.signal(s, lambda *_: handle_signal()) + + await stop_event.wait() + await nc.drain() + await nc.close() + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + pass diff --git a/files/nats_registration_listener.py b/files/nats_registration_listener.py index 740aeb2..ee4fb5e 100644 --- a/files/nats_registration_listener.py +++ b/files/nats_registration_listener.py @@ -24,7 +24,7 @@ import sys import hashlib import time from time import monotonic -from datetime import datetime, timezone +from datetime import datetime, timezone, timedelta from typing import Optional, Dict, Any, Tuple, Set from urllib.parse import urlencode from urllib.request import Request, urlopen @@ -62,6 +62,10 @@ RMQ_TIMEOUT = 3.0 # Example: 600000 = 10 minutes RMQ_DELAY_MS = 15000 +# Posture analyzer gate: skip re-running sot-updater-scheduler if recently run (seconds) +posture_analyzer = "sot-updater-scheduler" +sot_timeout = 1800 + # Product -> Tag slug mapping (kept unchanged, though not used now) PRODUCT_TAG_SLUG = { "fox100": "fox100-auto-upgrade-latest", @@ -225,16 +229,16 @@ async def nb_problem(log_status, msg: str): def nb_lookup_device_by_mac(mac: str, log_status) -> Tuple[ - Optional[str], Optional[int], Optional[int], Optional[str], Optional[Set[str]], Optional[Any], Optional[Any], Optional[Any], Optional[Any] + Optional[str], Optional[int], Optional[int], Optional[str], Optional[Set[str]], Optional[Any], Optional[Any], Optional[Any], Optional[Any], Optional[Any] ]: """ - Resolve MAC -> (hostname, iface_id, device_id, device_status_value, tag_slugs_set, action_next, action_last, action_next_timestamp, action_state) + Resolve MAC -> (hostname, iface_id, device_id, device_status_value, tag_slugs_set, action_next, action_last, action_next_timestamp, action_state, sot_ts) - Logs problems for anomalies (mac not found, unassigned, wrong type, iface fetch fail). - If device detail fetch fails, returns host/id with status/tags/custom_fields as None. """ mac_norm = normalize_mac(mac) if not mac_norm: - return None, None, None, None, None, None, None, None, None + return None, None, None, None, None, None, None, None, None, None # Step 0: short TTL cache (avoid repeated NetBox GETs for chatty devices) if NB_LOOKUP_CACHE_TTL > 0: @@ -243,13 +247,13 @@ def nb_lookup_device_by_mac(mac: str, log_status) -> Tuple[ if cached: exp, val = cached if exp > nowm: - host, iface_id, dev_id, status_val, tag_slugs, action_next, action_last, action_next_timestamp, action_state = val + host, iface_id, dev_id, status_val, tag_slugs, action_next, action_last, action_next_timestamp, action_state, sot_ts = val if tag_slugs is not None: try: tag_slugs = set(tag_slugs) except Exception: pass - return host, iface_id, dev_id, status_val, tag_slugs, action_next, action_last, action_next_timestamp, action_state + return host, iface_id, dev_id, status_val, tag_slugs, action_next, action_last, action_next_timestamp, action_state, sot_ts NB_LOOKUP_CACHE.pop(mac_norm, None) base = NB_URL.rstrip("/") @@ -269,12 +273,12 @@ def nb_lookup_device_by_mac(mac: str, log_status) -> Tuple[ data, code = http_get_json(f"{base}/api/dcim/mac-addresses/", params={"mac_address": mac_norm, "limit": "2"}, headers=h) if code != 200 or not data: asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: mac query http={code} mac={mac_norm}")) - return None, None, None, None, None, None, None, None, None + return None, None, None, None, None, None, None, None, None, None results = (data or {}).get("results") or [] if not results: asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: mac not found mac={mac_norm}")) - return None, None, None, None, None, None, None, None, None + return None, None, None, None, None, None, None, None, None, None rec = results[0] aot = (rec.get("assigned_object_type") or "").strip() @@ -285,10 +289,10 @@ def nb_lookup_device_by_mac(mac: str, log_status) -> Tuple[ if not aot or aoid is None: asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: mac unassigned mac={mac_norm}")) - return None, None, None, None, None, None, None, None, None + return None, None, None, None, None, None, None, None, None, None if aot != "dcim.interface": asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: mac assigned to {aot} mac={mac_norm} iface_id={aoid}")) - return None, aoid, None, None, None, None, None, None, None + return None, aoid, None, None, None, None, None, None, None, None # Step 2: Interface -> Device (shallow) iface, code2 = http_get_json(f"{base}/api/dcim/interfaces/{aoid}/", params={"fields": "device"}, headers=h) @@ -296,14 +300,14 @@ def nb_lookup_device_by_mac(mac: str, log_status) -> Tuple[ iface, code2 = http_get_json(f"{base}/api/dcim/interfaces/{aoid}/", headers=h) if code2 != 200 or not iface: asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: iface fetch http={code2} iface_id={aoid}")) - return None, aoid, None, None, None, None, None, None, None + return None, aoid, None, None, None, None, None, None, None, None dev = iface.get("device") or {} host = dev.get("name") or dev.get("display") dev_id = dev.get("id") if not host or dev_id is None: asyncio.create_task(nb_problem(log_status, f"[{ts()}] nb: iface has no device iface_id={aoid}")) - return None, aoid, None, None, None, None, None, None, None + return None, aoid, None, None, None, None, None, None, None, None # Step 3: Device detail (single fetch for status, tags, custom_fields.*) device, code3 = http_get_json(f"{base}/api/dcim/devices/{dev_id}/", params={"fields": "status,tags,custom_fields"}, headers=h) @@ -311,7 +315,7 @@ def nb_lookup_device_by_mac(mac: str, log_status) -> Tuple[ device, code3 = http_get_json(f"{base}/api/dcim/devices/{dev_id}/", headers=h) if code3 != 200 or not device: # treat as "no extra info" - return host, aoid, dev_id, None, None, None, None, None, None + return host, aoid, dev_id, None, None, None, None, None, None, None status_val = ((device.get("status") or {}).get("value")) or None tags = device.get("tags") or [] @@ -326,6 +330,7 @@ def nb_lookup_device_by_mac(mac: str, log_status) -> Tuple[ action_last = cf.get("action_last") action_next_timestamp = cf.get("action_next_timestamp") action_state = cf.get("action_state") + sot_ts = cf.get("sot_ts") # Step 4: populate cache (only on full success) if NB_LOOKUP_CACHE_TTL > 0: @@ -342,12 +347,13 @@ def nb_lookup_device_by_mac(mac: str, log_status) -> Tuple[ action_last, action_next_timestamp, action_state, + sot_ts, ), ) except Exception: pass - return host, aoid, dev_id, (status_val if isinstance(status_val, str) else None), tag_slugs, action_next, action_last, action_next_timestamp, action_state + return host, aoid, dev_id, (status_val if isinstance(status_val, str) else None), tag_slugs, action_next, action_last, action_next_timestamp, action_state, sot_ts # ========================= @@ -428,7 +434,7 @@ async def main(): if product == "fox100": nb_start = time.perf_counter() try: - host, iface_id, dev_id, status_val, tag_slugs, action_next, action_last, action_next_timestamp, action_state = nb_lookup_device_by_mac( + host, iface_id, dev_id, status_val, tag_slugs, action_next, action_last, action_next_timestamp, action_state, sot_ts = nb_lookup_device_by_mac( mac=mac, log_status=log_status ) if host: @@ -492,8 +498,30 @@ async def main(): now_epoch = int(time.time()) + # Gate: if action_next is posture_analyzer and sot_ts is recent, skip sending task + skip_due_sot = False + if action_next_str == posture_analyzer: + try: + if isinstance(sot_ts, str): + _st = sot_ts.strip() + if _st: + _dt = datetime.strptime(_st, "%d%m%y-%H%M%S").replace(tzinfo=timezone(timedelta(hours=2))) + _age = now_epoch - int(_dt.timestamp()) + if _age >= 0 and _age < sot_timeout: + skip_due_sot = True + async with print_lock: + print( + f"[{ts()}] skip action_next because sot_ts is recent host={host} task={action_next_str} age_s={_age} sot_ts={_st}", + file=sys.stdout, + flush=True, + ) + except Exception: + skip_due_sot = False + # If action_last == action_next, apply cooldown based on action_next_timestamp (600s) - if action_last_str == action_next_str: + if skip_due_sot: + pass + elif action_last_str == action_next_str: allow_repeat = True try: if action_next_timestamp is None: