#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
"""bulk-enrich.py -- enrich a CSV/log of IP addresses via the ip.xakki.pro API.

Reads a CSV with one IP address per row (default column name "ip"), calls
GET https://ip.xakki.pro/api/v1/ip?ip=<addr>&include=network,security for
each row, and writes a new CSV with extra columns: country, asn, as_name,
is_proxy, is_anonymous, is_malicious, enrich_status.

Zero dependencies -- stdlib only (urllib, json, csv, ipaddress, argparse,
time), no `pip install` required. Works on Python 3.8+.

Usage:
    python3 bulk-enrich.py ips.csv enriched.csv
    python3 bulk-enrich.py ips.csv enriched.csv --ip-column source_ip --delay 1.0

Rate-limit contract (read this before pointing it at a huge file):
  GET /api/v1/ip never returns a hard 429 from the app itself -- the "geo",
  "network" and "security" groups are metered PER UNIQUE IP PER DAY (plus a
  tighter per-minute budget), and going over the budget soft-withholds the
  group: HTTP 200, "status":"partial", the group's name in "limited", and
  RateLimit-Limit/RateLimit-Remaining/RateLimit-Reset headers on the
  response. A withheld group is recorded here as "withheld", never left
  blank/null -- a blank cell would silently look identical to "no data for
  this IP", which is a different fact. An anonymous caller's budget is quite
  small (see /api/v1/tiers or the RateLimit-Limit header for the live
  number -- it is not hardcoded here because it is admin-editable); register
  a free API token for a much higher one, or slow this script down with
  --delay if you see "withheld" often.
  A 429/503 CAN still happen at the edge (nginx's own per-IP flood guard,
  independent of the app-level metering above) if you hammer the endpoint
  too fast from one machine -- this script backs off exponentially on any
  429/5xx or transport error and honors a Retry-After header when present.
"""

import argparse
import csv
import ipaddress
import json
import sys
import time
import urllib.error
import urllib.request

DEFAULT_BASE_URL = "https://ip.xakki.pro"
DEFAULT_INCLUDE = "network,security"
USER_AGENT = "myip-bulk-enrich-recipe/1.0 (+https://ip.xakki.pro/tools/scripts/bulk-enrich.html)"
NEW_COLUMNS = [
    "country", "asn", "as_name", "is_proxy", "is_anonymous",
    "is_malicious", "enrich_status",
]


def is_valid_ip(value):
    """True if value parses as an IPv4 or IPv6 address (stdlib ipaddress)."""
    try:
        ipaddress.ip_address(value.strip())
        return True
    except ValueError:
        return False


def fetch_with_retry(url, max_retries, base_backoff):
    """GET url, retrying on 429/5xx and transport errors with exponential
    backoff. Honors a Retry-After header (seconds) when the server sends
    one, otherwise doubles base_backoff each attempt. Returns the parsed
    JSON body, or raises the last exception after exhausting retries."""
    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT, "Accept": "application/json"})
    attempt = 0
    while True:
        try:
            with urllib.request.urlopen(req, timeout=15) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as e:
            retryable = e.code == 429 or e.code >= 500
            if not retryable or attempt >= max_retries:
                raise
            wait = float(e.headers.get("Retry-After", 0)) or (base_backoff * (2 ** attempt))
        except (urllib.error.URLError, TimeoutError, ConnectionError) as e:
            if attempt >= max_retries:
                raise
            wait = base_backoff * (2 ** attempt)
        attempt += 1
        time.sleep(wait)


def enrich_row(base_url, include, ip_value, max_retries, base_backoff):
    """Returns a dict of NEW_COLUMNS for one IP, never raising -- every
    failure mode (invalid input, transport error, unexpected response)
    degrades to an explicit status string instead of a blank/null cell."""
    out = {c: "" for c in NEW_COLUMNS}
    if not is_valid_ip(ip_value):
        out["enrich_status"] = "invalid_ip"
        return out

    url = "{}/api/v1/ip?ip={}&include={}".format(base_url, ip_value.strip(), include)
    try:
        data = fetch_with_retry(url, max_retries, base_backoff)
    except Exception as e:  # noqa: BLE001 -- one recipe-wide catch, logged below
        out["enrich_status"] = "error: {}".format(e)
        return out

    # Guard against the API's own documented behavior: an invalid `ip` query
    # param is silently IGNORED (not rejected), and the endpoint falls back
    # to answering about the CALLER's own IP instead. We already validated
    # ip_value above, so this should never fire -- but if it ever does, never
    # attribute the wrong IP's data to this row.
    try:
        if ipaddress.ip_address(data.get("ip", "")) != ipaddress.ip_address(ip_value.strip()):
            out["enrich_status"] = "unexpected_response: ip mismatch"
            return out
    except ValueError:
        out["enrich_status"] = "unexpected_response: no ip in response"
        return out

    limited = set(data.get("limited") or [])
    network = data.get("network") or {}
    geo = data.get("geo") or {}
    privacy = data.get("privacy")
    threat = data.get("threat")

    out["country"] = geo.get("countryCode", "") or ""
    out["asn"] = network.get("asn", "") or ""
    out["as_name"] = network.get("asName", "") or ""

    if "security" in limited:
        out["is_proxy"] = "withheld"
        out["is_anonymous"] = "withheld"
        out["is_malicious"] = "withheld"
    else:
        out["is_proxy"] = str(privacy.get("isProxy")) if privacy else ""
        out["is_anonymous"] = str(privacy.get("isAnonymous")) if privacy else ""
        out["is_malicious"] = str(threat.get("isMalicious")) if threat else ""

    out["enrich_status"] = "partial" if data.get("status") == "partial" else "ok"
    return out


def main():
    parser = argparse.ArgumentParser(description="Enrich a CSV of IP addresses via the ip.xakki.pro API.")
    parser.add_argument("input_csv", help="Path to the input CSV (must have an IP column).")
    parser.add_argument("output_csv", help="Path to write the enriched CSV to.")
    parser.add_argument("--ip-column", default="ip", help='Column holding the IP address (default: "ip").')
    parser.add_argument("--include", default=DEFAULT_INCLUDE, help="Groups to request (default: network,security).")
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL, help="API base URL (default: production).")
    parser.add_argument("--delay", type=float, default=0.5, help="Seconds to sleep between requests (default: 0.5).")
    parser.add_argument("--max-retries", type=int, default=5, help="Max retries per request on 429/5xx (default: 5).")
    args = parser.parse_args()

    with open(args.input_csv, newline="", encoding="utf-8") as f_in:
        reader = csv.DictReader(f_in)
        if args.ip_column not in (reader.fieldnames or []):
            sys.exit('error: column "{}" not found in {} (columns: {})'.format(
                args.ip_column, args.input_csv, reader.fieldnames))
        fieldnames = list(reader.fieldnames) + [c for c in NEW_COLUMNS if c not in reader.fieldnames]
        rows = list(reader)

    counts = {"ok": 0, "partial": 0, "invalid_ip": 0, "error": 0}
    with open(args.output_csv, "w", newline="", encoding="utf-8") as f_out:
        writer = csv.DictWriter(f_out, fieldnames=fieldnames)
        writer.writeheader()
        for i, row in enumerate(rows):
            enrichment = enrich_row(args.base_url, args.include, row.get(args.ip_column, ""),
                                     args.max_retries, args.delay)
            row.update(enrichment)
            writer.writerow(row)
            status = enrichment["enrich_status"]
            key = status if status in counts else ("error" if status.startswith("error") else status)
            counts[key] = counts.get(key, 0) + 1
            print("[{}/{}] {} -> {}".format(i + 1, len(rows), row.get(args.ip_column, ""), status), file=sys.stderr)
            if i < len(rows) - 1:
                time.sleep(args.delay)

    print("done: {} written to {}".format(counts, args.output_csv), file=sys.stderr)


if __name__ == "__main__":
    main()