<?php

/**
 * GeoLite2ASNExtractor
 *
 * Pure-PHP reader for MaxMind's GeoLite2-ASN.mmdb that walks the entire
 * binary search tree and returns every network grouped by ASN, e.g.:
 *
 *   [
 *     "as1400" => ["1.0.0.0/32", "1.0.0.1/32", ...],
 *     "as13335" => ["1.1.1.0/24", ...],
 *   ]
 *
 * No Composer dependency required — implements the MaxMind DB binary
 * format (search tree + data section) directly per the public spec:
 * https://maxmind.github.io/MaxMind-DB/
 *
 * Usage:
 *   $extractor = new GeoLite2ASNExtractor('/path/to/GeoLite2-ASN.mmdb');
 *   $asnMap = $extractor->extract();
 *
 * Notes:
 *   - GeoLite2-ASN.mmdb is normally built as an IPv6 tree with the IPv4
 *     address space embedded under ::/96. This class detects that and
 *     emits plain IPv4 CIDRs (e.g. "1.0.0.0/32"), skipping true
 *     IPv6-only ranges. Set $includeIPv6 = true in extract() if you
 *     want IPv6 ranges included too (returned as e.g. "2606:4700::/32").
 *   - The whole file is loaded into memory once (file_get_contents).
 *     GeoLite2-ASN.mmdb is tens of MB, which is fine for most setups,
 *     but be aware if running under a memory-constrained environment.
 *   - Walking the full tree is O(number of tree nodes), which can be a
 *     few million recursive calls for the ASN database. This can take
 *     anywhere from several seconds to a couple of minutes depending on
 *     hardware. Consider raising max_execution_time for CLI/cron use.
 */
class GeoLite2ASNExtractor
{
    private const METADATA_MARKER = "\xab\xcd\xefMaxMind.com";

    private string $data;
    private array $metadata = [];

    private int $nodeCount = 0;
    private int $recordSize = 0;
    private int $nodeByteSize = 0;
    private int $searchTreeSize = 0;
    private int $dataSectionStart = 0;
    private int $ipVersion = 6;

    public function __construct(string $mmdbPath)
    {
        if (!is_readable($mmdbPath)) {
            throw new \RuntimeException("Cannot read file: {$mmdbPath}");
        }

        $contents = file_get_contents($mmdbPath);
        if ($contents === false) {
            throw new \RuntimeException("Failed to read file: {$mmdbPath}");
        }

        $this->data = $contents;
        $this->parseMetadata();
    }

    /**
     * Optional convenience helper: download an .mmdb file over HTTP(S)
     * and construct an extractor from it.
     */
    public static function fromUrl(string $url, ?string $saveTo = null): self
    {
        $tmpPath = $saveTo ?? tempnam(sys_get_temp_dir(), 'mmdb_');

        $ch = curl_init($url);
        $fh = fopen($tmpPath, 'wb');
        if ($fh === false) {
            throw new \RuntimeException("Cannot open temp file for writing: {$tmpPath}");
        }

        curl_setopt_array($ch, [
            CURLOPT_FILE => $fh,
            CURLOPT_FOLLOWLOCATION => true,
            CURLOPT_MAXREDIRS => 5,
            CURLOPT_TIMEOUT => 300,
            CURLOPT_FAILONERROR => true,
            CURLOPT_USERAGENT => 'GeoLite2ASNExtractor/1.0',
        ]);

        $ok = curl_exec($ch);
        $err = curl_error($ch);
        curl_close($ch);
        fclose($fh);

        if ($ok === false) {
            if ($saveTo === null) {
                @unlink($tmpPath);
            }
            throw new \RuntimeException("Download failed: {$err}");
        }

        return new self($tmpPath);
    }

    /**
     * Walk the entire tree and return ["as<N>" => ["a.b.c.d/x", ...], ...]
     */
    public function extract(bool $includeIPv6 = false): array
    {
        $result = [];
        $this->walk(0, '', $result, $includeIPv6);
        return $result;
    }

    public function getMetadata(): array
    {
        return $this->metadata;
    }

    // ------------------------------------------------------------------
    // Search tree traversal
    // ------------------------------------------------------------------

    /**
     * @param string $bits '0'/'1' characters representing the path taken so far
     */
    private function walk(int $recordValue, string $bits, array &$result, bool $includeIPv6): void
    {
        if ($recordValue === $this->nodeCount) {
            // Empty branch, no data assigned here.
            return;
        }

        if ($recordValue > $this->nodeCount) {
            // Leaf: record points into the data section.
            //
            // Per the MaxMind DB spec, the offset within the data section is
            // (record_value - node_count - 16), and the data section itself
            // starts at (search_tree_size + 16) — the two 16s cancel out, so
            // the correct absolute file offset is search_tree_size plus
            // (record_value - node_count). Using data_section_start here
            // instead would shift every leaf read by 16 bytes.
            $absoluteOffset = $this->searchTreeSize + ($recordValue - $this->nodeCount);
            [$record, ] = $this->decodeData($absoluteOffset);
            $this->addResult($bits, $record, $result, $includeIPv6);
            return;
        }

        // Otherwise it's an index of another node — recurse both children.
        [$left, $right] = $this->readNode($recordValue);
        $this->walk($left, $bits . '0', $result, $includeIPv6);
        $this->walk($right, $bits . '1', $result, $includeIPv6);
    }

    private function readNode(int $nodeNumber): array
    {
        $offset = $nodeNumber * $this->nodeByteSize;
        $bytes = substr($this->data, $offset, $this->nodeByteSize);

        if ($this->recordSize === 24) {
            $left = (ord($bytes[0]) << 16) | (ord($bytes[1]) << 8) | ord($bytes[2]);
            $right = (ord($bytes[3]) << 16) | (ord($bytes[4]) << 8) | ord($bytes[5]);
        } elseif ($this->recordSize === 28) {
            $middle = ord($bytes[3]);
            $left = (ord($bytes[0]) << 16) | (ord($bytes[1]) << 8) | ord($bytes[2]);
            $left |= ($middle >> 4) << 24;
            $right = (($middle & 0x0f) << 24) | (ord($bytes[4]) << 16) | (ord($bytes[5]) << 8) | ord($bytes[6]);
        } elseif ($this->recordSize === 32) {
            $left = (ord($bytes[0]) << 24) | (ord($bytes[1]) << 16) | (ord($bytes[2]) << 8) | ord($bytes[3]);
            $right = (ord($bytes[4]) << 24) | (ord($bytes[5]) << 16) | (ord($bytes[6]) << 8) | ord($bytes[7]);
        } else {
            throw new \RuntimeException("Unsupported record_size: {$this->recordSize}");
        }

        return [$left, $right];
    }

    private function addResult(string $bits, $record, array &$result, bool $includeIPv6): void
    {
        if (!is_array($record) || !isset($record['autonomous_system_number'])) {
            return; // no ASN data at this leaf
        }

        $asn = 'as' . $record['autonomous_system_number'];
        $prefixLen = strlen($bits);

        if ($this->ipVersion === 6) {
            if ($prefixLen >= 96 && substr($bits, 0, 96) === str_repeat('0', 96)) {
                // Embedded IPv4 space (::/96).
                $ipv4Bits = substr($bits, 96);
                $cidr = $this->bitsToIPv4($ipv4Bits) . '/' . strlen($ipv4Bits);
                $result[$asn][] = $cidr;
            } elseif ($includeIPv6) {
                $cidr = $this->bitsToIPv6($bits) . '/' . $prefixLen;
                $result[$asn][] = $cidr;
            }
            // else: real IPv6 range but caller doesn't want IPv6 — skip.
        } else {
            // Pure IPv4 tree (older / rare db builds).
            $cidr = $this->bitsToIPv4($bits) . '/' . $prefixLen;
            $result[$asn][] = $cidr;
        }
    }

    private function bitsToIPv4(string $bits): string
    {
        $bits = str_pad($bits, 32, '0'); // pad host bits with 0 -> network address
        $octets = [];
        for ($i = 0; $i < 4; $i++) {
            $octets[] = bindec(substr($bits, $i * 8, 8));
        }
        return implode('.', $octets);
    }

    private function bitsToIPv6(string $bits): string
    {
        $bits = str_pad($bits, 128, '0');
        $groups = [];
        for ($i = 0; $i < 8; $i++) {
            $groups[] = dechex(bindec(substr($bits, $i * 16, 16)));
        }
        $expanded = implode(':', $groups);
        return inet_ntop(inet_pton($expanded)) ?: $expanded;
    }

    // ------------------------------------------------------------------
    // Metadata / data section decoding (MaxMind DB data format)
    // ------------------------------------------------------------------

    private function parseMetadata(): void
    {
        $pos = strrpos($this->data, self::METADATA_MARKER);
        if ($pos === false) {
            throw new \RuntimeException('MaxMind DB metadata marker not found — not a valid .mmdb file');
        }

        $offset = $pos + strlen(self::METADATA_MARKER);
        [$metadata, ] = $this->decodeData($offset);

        if (!is_array($metadata) || !isset($metadata['node_count'], $metadata['record_size'], $metadata['ip_version'])) {
            throw new \RuntimeException('Malformed MaxMind DB metadata');
        }

        $this->metadata = $metadata;
        $this->nodeCount = (int) $metadata['node_count'];
        $this->recordSize = (int) $metadata['record_size'];
        $this->ipVersion = (int) $metadata['ip_version'];
        $this->nodeByteSize = (int) (($this->recordSize * 2) / 8);
        $this->searchTreeSize = $this->nodeCount * $this->nodeByteSize;
        $this->dataSectionStart = $this->searchTreeSize + 16; // 16-byte all-zero separator
    }

    /**
     * Decodes one MaxMind DB data item starting at $offset.
     * Returns [decodedValue, offsetAfterThisItem].
     */
    private function decodeData(int $offset): array
    {
        $ctrl = ord($this->data[$offset]);
        $offset++;

        $type = $ctrl >> 5;   // top 3 bits
        $size = $ctrl & 0x1f; // bottom 5 bits

        if ($type === 0) {
            // Extended type: real type = next byte + 7
            $type = 7 + ord($this->data[$offset]);
            $offset++;
        }

        if ($type === 1) {
            return $this->decodePointer($ctrl, $offset);
        }

        if ($type !== 14) { // not boolean — booleans store their value in $size directly
            if ($size === 29) {
                $size = 29 + ord($this->data[$offset]);
                $offset += 1;
            } elseif ($size === 30) {
                $size = 285 + $this->readUint($offset, 2);
                $offset += 2;
            } elseif ($size === 31) {
                $size = 65821 + $this->readUint($offset, 3);
                $offset += 3;
            }
        }

        switch ($type) {
            case 2: // utf8_string
            case 4: // bytes
                $value = substr($this->data, $offset, $size);
                $offset += $size;
                return [$value, $offset];

            case 3: // double
                $bytes = substr($this->data, $offset, 8);
                $offset += 8;
                $u = unpack('E', $bytes);
                return [$u[1], $offset];

            case 5: // uint16
            case 6: // uint32
            case 9: // uint64 (may lose precision above 2^63 — not expected for ASN fields)
                $value = $this->readUint($offset, $size);
                $offset += $size;
                return [$value, $offset];

            case 7: // map
                $map = [];
                for ($i = 0; $i < $size; $i++) {
                    [$key, $offset] = $this->decodeData($offset);
                    [$val, $offset] = $this->decodeData($offset);
                    $map[$key] = $val;
                }
                return [$map, $offset];

            case 8: // int32
                $value = $this->readUint($offset, $size);
                if ($size > 0 && ($value & (1 << (8 * $size - 1)))) {
                    $value -= (1 << (8 * $size));
                }
                $offset += $size;
                return [$value, $offset];

            case 10: // uint128 — returned as a hex string (not used by ASN records)
                $hex = bin2hex(substr($this->data, $offset, $size));
                $offset += $size;
                return [$hex, $offset];

            case 11: // array
                $arr = [];
                for ($i = 0; $i < $size; $i++) {
                    [$val, $offset] = $this->decodeData($offset);
                    $arr[] = $val;
                }
                return [$arr, $offset];

            case 13: // end marker
                return [null, $offset];

            case 14: // boolean — value is $size itself (0 or 1)
                return [$size === 1, $offset];

            case 15: // float
                $bytes = substr($this->data, $offset, 4);
                $offset += 4;
                $u = unpack('G', $bytes);
                return [$u[1], $offset];

            default:
                throw new \RuntimeException("Unsupported MaxMind DB data type: {$type} at offset {$offset}");
        }
    }

    private function decodePointer(int $ctrl, int $offset): array
    {
        $ptrSize = ($ctrl >> 3) & 0x03;
        $valueHigh = $ctrl & 0x07;

        switch ($ptrSize) {
            case 0:
                $pointer = ($valueHigh << 8) | ord($this->data[$offset]);
                $offset += 1;
                $base = 0;
                break;
            case 1:
                $pointer = ($valueHigh << 16) | $this->readUint($offset, 2);
                $offset += 2;
                $base = 2048;
                break;
            case 2:
                $pointer = ($valueHigh << 24) | $this->readUint($offset, 3);
                $offset += 3;
                $base = 526336;
                break;
            default: // 3
                $pointer = $this->readUint($offset, 4);
                $offset += 4;
                $base = 0;
                break;
        }

        $target = $this->dataSectionStart + $pointer + $base;
        [$value, ] = $this->decodeData($target);

        return [$value, $offset];
    }

    private function readUint(int $offset, int $numBytes): int
    {
        $value = 0;
        for ($i = 0; $i < $numBytes; $i++) {
            $value = ($value << 8) | ord($this->data[$offset + $i]);
        }
        return $value;
    }
}
