riposte-marketplace/integrations/harfanglab/scripts/endpoint_count.py

import json, os, re, sys, time, urllib.request, urllib.parse, urllib.error


RETRIES = 3


def retryable(method, status):
    # HarfangLab answers a 500 now and then on a call that passes when sent again.
    # A read is safe to resend; a write only on an answer that says the action was
    # not run (rate limit, gateway down), so an isolation or a job is never sent twice.
    if status in (429, 502, 503):
        return True
    return method == "GET" and status in (500, 504)


def request(method, url, headers, body=None):
    data = json.dumps(body).encode("utf-8") if body is not None else None
    for attempt in range(RETRIES + 1):
        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                raw = resp.read()
                return json.loads(raw) if raw else {}
        except urllib.error.HTTPError as e:
            if attempt == RETRIES or not retryable(method, e.code):
                raise
            e.close()
            time.sleep(2 ** attempt)


def qs(params):
    clean = {}
    for k, v in params.items():
        if v is None or v == "" or v == []:
            continue
        if isinstance(v, bool):
            v = "true" if v else "false"
        clean[k] = v
    return ("?" + urllib.parse.urlencode(clean)) if clean else ""


def csv(v):
    return [x.strip() for x in str(v or "").split(",") if x.strip()]


# Agent fields HarfangLab exposes as exact-match query filters on
# /api/data/endpoint/Agent/. Input name == API filter name.
EXACT_FILTERS = [
    "hostname",
    "domainname",
    "domain",
    "ipaddress",
    "external_ipaddress",
    "ostype",
    "osproducttype",
    "osversion",
    "producttype",
    "version",
    "status",
    "isolation_state",
    "tenant",
]

# Convenience bounds mapped onto Django REST Framework range lookups.
RANGE_FILTERS = {
    "lastseen_after": "lastseen__gte",
    "lastseen_before": "lastseen__lte",
    "firstseen_after": "firstseen__gte",
    "firstseen_before": "firstseen__lte",
}


# Version bounds. The API matches a version only exactly, and a range lookup it
# does not support would be dropped without a word — so these are compared here,
# number by number (4.10 is above 4.9, which a string comparison denies).
VERSION_BOUNDS = {
    "version_lt": lambda v, t: v < t,
    "version_lte": lambda v, t: v <= t,
    "version_gt": lambda v, t: v > t,
    "version_gte": lambda v, t: v >= t,
}
PAGE = 500


def parse_version(v):
    parts = re.findall(r"\d+", str(v or ""))
    return tuple(int(x) for x in parts) if parts else None


def pad(a, b):
    n = max(len(a), len(b))
    return a + (0,) * (n - len(a)), b + (0,) * (n - len(b))


def to_bool(v):
    if isinstance(v, bool):
        return v
    return str(v).strip().lower() in ("true", "1", "yes")


def main():
    secrets = json.loads(os.environ.get("INTEGRATION_SECRETS", "{}"))
    inputs = json.loads(os.environ.get("INTEGRATION_INPUTS", "{}"))
    base = secrets.get("url", "").rstrip("/")
    headers = {
        "Authorization": "Token " + secrets.get("api_token", ""),
        "Accept": "application/json",
        "Content-Type": "application/json",
    }
    # === per-command logic ===
    # An instance bound to one tenant of a shared, multi-tenant manager carries it
    # in its config: it scopes the count, and the fleet total beside it, unless the
    # block names a tenant itself.
    scope = {}
    if not str(inputs.get("tenant") or "").strip() and str(secrets.get("tenant") or "").strip():
        scope["tenant"] = str(secrets["tenant"]).strip()
    filters = dict(scope)
    for field in EXACT_FILTERS:
        value = inputs.get(field)
        if value is None or value == "":
            continue
        filters[field] = to_bool(value) if field == "isolation_state" else value
    for field, lookup in RANGE_FILTERS.items():
        value = inputs.get(field)
        if value:
            filters[lookup] = value
    # Escape hatch for any filter this command does not name explicitly,
    # e.g. "groups__name=Serveurs&osmajor=10".
    for tok in str(inputs.get("filters") or "").split("&"):
        kv = tok.split("=", 1)
        if len(kv) == 2 and kv[0].strip():
            filters[kv[0].strip()] = kv[1]

    bounds = {}
    for field in VERSION_BOUNDS:
        value = str(inputs.get(field) or "").strip()
        if value:
            target = parse_version(value)
            if target is None:
                raise ValueError(field + " must be a version such as 4.10.2, got: " + repr(value))
            bounds[field] = target

    # The agent list is a DRF paginated collection: `count` is the total number
    # of matches, independent of the page size. limit=1 keeps the payload tiny.
    agents = base + "/api/data/endpoint/Agent/"
    if bounds:
        # A version bound reads every agent the other filters keep and counts
        # here. An agent with no readable version is left out and counted apart,
        # rather than guessed above or below the bound.
        count, scanned, unreadable, offset = 0, 0, 0, 0
        while True:
            page = request("GET", agents + qs(dict(filters, limit=PAGE, offset=offset)), headers)
            rows = page.get("results") or []
            for agent in rows:
                scanned += 1
                version = parse_version(agent.get("version"))
                if version is None:
                    unreadable += 1
                    continue
                if all(VERSION_BOUNDS[f](*pad(version, t)) for f, t in bounds.items()):
                    count += 1
            offset += len(rows)
            if not rows or offset >= page.get("count", 0):
                break
        result = {"count": count, "filters": dict(filters, **{f: inputs[f] for f in bounds}),
                  "scanned": scanned, "unreadable_version": unreadable}
    else:
        matched = request("GET", agents + qs(dict(filters, limit=1, offset=0)), headers)
        result = {"count": matched.get("count", 0), "filters": filters}

    # The API silently drops a query parameter it does not recognise as a
    # filter, which yields the unfiltered total dressed up as a filtered count.
    # Returning the fleet total alongside makes that failure visible.
    if filters != scope or bounds:
        everything = request("GET", agents + qs(dict(scope, limit=1, offset=0)), headers)
        result["total"] = everything.get("count", 0)
    else:
        result["total"] = result["count"]

    print(json.dumps(result))


try:
    main()
except urllib.error.HTTPError as e:
    print(json.dumps({"error": "HTTP " + str(e.code), "detail": e.read().decode("utf-8", "replace")}))
    sys.exit(1)
except Exception as e:
    print(json.dumps({"error": str(e)}))
    sys.exit(1)