#!/usr/bin/env python3
# SPDX-License-Identifier: MIT
"""asn-lookup.py -- a single asn_lookup(ip) function wrapping
GET /api/v1/ip?ip=<addr>&include=network, with clean error handling.

Shown alongside asn-lookup.js -- the same task in Python and JavaScript/Node,
the two languages this section's audience is most likely to already have
installed, so you can copy whichever one matches your stack without
translating the other by hand.

Zero dependencies -- stdlib only (urllib, json, ipaddress). Python 3.8+.

Usage as a library:
    from importlib.machinery import SourceFileLoader
    m = SourceFileLoader("asn_lookup", "asn-lookup.py").load_module()
    print(m.asn_lookup("8.8.8.8"))

Usage as a CLI:
    python3 asn-lookup.py 8.8.8.8
"""

import ipaddress
import json
import sys
import urllib.error
import urllib.request

DEFAULT_BASE_URL = "https://ip.xakki.pro"
USER_AGENT = "myip-asn-lookup-recipe/1.0 (+https://ip.xakki.pro/tools/scripts/asn-lookup.html)"


class InvalidIPError(ValueError):
    """Raised when the input string is not a valid IPv4/IPv6 address."""


class RateLimitedError(RuntimeError):
    """Raised when the API's "network" group came back rate-limited
    (soft-withheld -- see the rate-limit note in bulk-enrich.py) for THIS
    lookup, so no ASN data is available right now."""


def asn_lookup(ip, base_url=DEFAULT_BASE_URL, timeout=15):
    """Returns a dict {asn, as_name, prefix, registry, allocated, country,
    classification_type} for ip. Raises InvalidIPError for a malformed
    input (checked BEFORE any request -- an invalid `ip` query param is
    silently ignored by the API and answers about the caller's own IP
    instead, which would silently mislabel the result), RateLimitedError
    when the network group is withheld this call, or urllib.error.URLError
    / json.JSONDecodeError for a genuine transport/protocol failure."""
    try:
        ipaddress.ip_address(ip)
    except ValueError:
        raise InvalidIPError('"{}" is not a valid IPv4 or IPv6 address'.format(ip))

    url = "{}/api/v1/ip?ip={}&include=network".format(base_url, ip)
    req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT, "Accept": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        data = json.loads(resp.read().decode("utf-8"))

    if ipaddress.ip_address(data.get("ip", ip)) != ipaddress.ip_address(ip):
        raise RuntimeError("unexpected response: ip mismatch (got {!r})".format(data.get("ip")))

    if "network" in (data.get("limited") or []):
        raise RateLimitedError("the network group was rate-limited for this call; retry later")

    network = data.get("network") or {}
    rdap = network.get("rdap") or {}
    classification = network.get("classification") or {}
    return {
        "asn": network.get("asn"),
        "as_name": network.get("asName"),
        "prefix": network.get("prefix"),
        "registry": network.get("registry"),
        "allocated": network.get("allocated"),
        "country": rdap.get("country"),
        "classification_type": classification.get("type") or None,
    }


def main():
    if len(sys.argv) != 2:
        sys.exit("usage: python3 asn-lookup.py <ip>")
    ip = sys.argv[1]
    try:
        result = asn_lookup(ip)
    except InvalidIPError as e:
        sys.exit("invalid input: {}".format(e))
    except RateLimitedError as e:
        sys.exit("rate limited: {}".format(e))
    except (urllib.error.URLError, json.JSONDecodeError) as e:
        sys.exit("request failed: {}".format(e))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()