#!/usr/bin/env python3
"""Stage verified Data-Shield feeds for CrowdSec or nftables. GPL-3.0-only.

Python 3.11+. This helper never executes commands or changes a firewall or API.
The caller supplies a new output directory inside a trusted private parent.
"""
import argparse
import hashlib
import ipaddress
import json
import re
import shutil
import sys
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path

SNAPSHOT = 'https://data-shield.eu/data/latest.json'
MAIN = 'https://api.github.com/repos/duggytuxy/Data-Shield_IPv4_Blocklist/git/ref/heads/main'
RAW = 'https://raw.githubusercontent.com/duggytuxy/Data-Shield_IPv4_Blocklist/'
FILES = {'standard': 'prod_data-shield_ipv4_blocklist.txt',
         'critical': 'prod_critical_data-shield_ipv4_blocklist.txt'}
MAX_BYTES = 12 * 1024 * 1024


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, request, fp, code, message, headers, url):
        raise ValueError('Unexpected HTTP redirect')


def fetch(url):
    request = urllib.request.Request(url, headers={'User-Agent': 'Data-Shield-Linux/1.0',
                                                  'Accept-Encoding': 'identity'})
    with urllib.request.build_opener(NoRedirect()).open(request, timeout=60) as response:
        if response.status != 200 or response.url != url:
            raise ValueError('Unexpected HTTP response')
        raw = response.read(MAX_BYTES + 1)
    if len(raw) > MAX_BYTES:
        raise ValueError('Response exceeds size limit')
    return raw


def commit(raw):
    value = json.loads(raw)['object']['sha']
    if not isinstance(value, str) or not re.fullmatch(r'[a-f0-9]{40}', value):
        raise ValueError('Invalid GitHub commit')
    return value


def metadata(raw, now=None):
    data = json.loads(raw)
    if data.get('schema_version') != 1:
        raise ValueError('Unsupported snapshot schema')
    generation = data['generation']
    if not re.fullmatch(r'[a-f0-9]{40}', generation['source_commit']):
        raise ValueError('Invalid snapshot commit')
    now = now or datetime.now(timezone.utc)
    created = datetime.fromisoformat(generation['created_at'].replace('Z', '+00:00'))
    if created.tzinfo is None or not now - timedelta(hours=6) <= created <= now + timedelta(minutes=5):
        raise ValueError('Publication is stale or has an invalid timestamp')
    profiles = data['profiles']
    for profile in FILES:
        entry = profiles[profile]
        for key, limit in (('blocked', 1000000), ('bytes', MAX_BYTES)):
            if type(entry[key]) is not int or not 1 <= entry[key] <= limit:
                raise ValueError('Invalid publication count or size')
        if not re.fullmatch(r'[a-f0-9]{64}', entry['sha256']):
            raise ValueError('Invalid publication digest')
    return generation, {p: {k: profiles[p][k] for k in ('blocked', 'bytes', 'sha256')} for p in FILES}


def addresses(raw):
    if not raw or len(raw) > MAX_BYTES or not raw.endswith(b'\n'):
        raise ValueError('Expected a complete nonempty LF-terminated IPv4 feed')
    result, previous = [], -1
    for line in raw[:-1].split(b'\n'):
        if not re.fullmatch(rb'[0-9.]+', line):
            raise ValueError('Unexpected content in IPv4 feed')
        address = ipaddress.IPv4Address(line.decode('ascii'))
        if str(address).encode() != line or not address.is_global or address.is_multicast or address.is_reserved:
            raise ValueError('Noncanonical or non-public IPv4')
        if int(address) <= previous:
            raise ValueError('Duplicate or unordered IPv4 feed')
        previous = int(address)
        result.append(str(address))
    return result


def publication(downloader=fetch):
    head = commit(downloader(MAIN))
    generation, expected = metadata(downloader(SNAPSHOT))
    bodies, members = {}, {}
    for profile, filename in FILES.items():
        raw = downloader(RAW + head + '/' + filename)
        entry = expected[profile]
        if len(raw) != entry['bytes'] or hashlib.sha256(raw).hexdigest() != entry['sha256']:
            raise ValueError(profile + ': SHA-256 or byte count differs from snapshot; retry later')
        members[profile] = addresses(raw)
        if len(members[profile]) != entry['blocked']:
            raise ValueError(profile + ': record count mismatch')
        bodies[profile] = raw
    after, after_profiles = metadata(downloader(SNAPSHOT))
    if head != commit(downloader(MAIN)) or generation != after or expected != after_profiles:
        raise ValueError('Publication changed during retrieval; retry later')
    manifest = {'verified_at': datetime.now(timezone.utc).isoformat(), 'source_commit': head,
                'snapshot_generation': generation, 'profiles': expected}
    return manifest, bodies, members


def transaction(members):
    lines = ['# Verified IPv4 members. Only dedicated Data-Shield sets are replaced.']
    for profile in FILES:
        lines.append('flush set inet data_shield ' + profile)
        values = members[profile]
        for start in range(0, len(values), 1000):
            lines.append('add element inet data_shield ' + profile + ' { ' + ', '.join(values[start:start + 1000]) + ' }')
    return '\n'.join(lines) + '\n'


def stage(output, downloader=fetch):
    # Validate both complete bodies before creating any output.
    manifest, bodies, members = publication(downloader)
    output.mkdir(mode=0o700, parents=False, exist_ok=False)
    try:
        for profile, raw in bodies.items():
            (output / (profile + '.txt')).write_bytes(raw)
        (output / 'sets.nft').write_text(transaction(members))
        (output / 'manifest.json').write_text(json.dumps(manifest, indent=2) + '\n')
    except BaseException:
        shutil.rmtree(output)
        raise
    return manifest


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', required=True, type=Path, help='New directory inside a trusted private parent')
    args = parser.parse_args()
    try:
        manifest = stage(args.output)
    except (OSError, ValueError, KeyError, TypeError) as error:
        print('Validation failed: ' + str(error), file=sys.stderr)
        return 1
    print(json.dumps(manifest, sort_keys=True))
    return 0


if __name__ == '__main__':
    sys.exit(main())