#!/usr/bin/env python3
"""Validate both public feeds and optionally update dedicated RouterOS v7 lists.

Run on a Linux management host with Python 3.11+ and OpenSSH, not on RouterOS.
No router connection is made without --apply. License: GPL-3.0-only.
"""
import argparse
import fcntl
import hashlib
import ipaddress
import json
import os
import re
import subprocess
import tempfile
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path

PROFILES = ('standard', 'critical')
ORIGIN = 'https://raw.githubusercontent.com/duggytuxy/Data-Shield_IPv4_Blocklist/refs/heads/main/'
URLS = {
    'standard': ORIGIN + 'prod_data-shield_ipv4_blocklist.txt',
    'critical': ORIGIN + 'prod_critical_data-shield_ipv4_blocklist.txt',
}
SNAPSHOT = 'https://data-shield.eu/data/latest.json'
MAX_BYTES = 12 * 1024 * 1024
MARKER = 'Data-Shield managed'


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-RouterOS/1.0'})
    with urllib.request.build_opener(NoRedirect()).open(request, timeout=60) as response:
        if response.status != 200 or response.url != url:
            raise ValueError('Unexpected HTTP status or redirect')
        raw = response.read(MAX_BYTES + 1)
    if len(raw) > MAX_BYTES:
        raise ValueError('Response exceeds the complete-file size limit')
    return raw


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 publication 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=24) <= created <= now + timedelta(minutes=5):
        raise ValueError('Publication is stale or has an invalid timestamp')
    profiles = data['profiles']
    for profile in PROFILES:
        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 PROFILES}


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):
    generation, expected = metadata(downloader(SNAPSHOT))
    result = {}
    for profile in PROFILES:
        raw = downloader(URLS[profile])
        entry = expected[profile]
        if len(raw) != entry['bytes'] or hashlib.sha256(raw).hexdigest() != entry['sha256']:
            raise ValueError(profile + ': full-file SHA-256 or byte count differs from the snapshot; retry later')
        result[profile] = addresses(raw)
        if len(result[profile]) != entry['blocked']:
            raise ValueError(profile + ': record count mismatch')
    after, after_profiles = metadata(downloader(SNAPSHOT))
    if generation != after or expected != after_profiles:
        raise ValueError('Publication changed during retrieval; retry later')
    return generation, expected, result


def list_name(profile, slot):
    if profile not in PROFILES or slot not in ('A', 'B'):
        raise ValueError('Invalid dedicated list name')
    return 'DS_' + profile.upper() + '_' + slot


def rule(profile):
    return '[/ip firewall filter find where comment="Data-Shield ' + profile.title() + '"]'


class Router:
    def __init__(self, alias):
        if not re.fullmatch(r'[A-Za-z0-9][A-Za-z0-9_.-]{0,63}', alias):
            raise ValueError('Use a simple SSH configuration alias, without a username or shell syntax')
        self.alias = alias
        self.options = ['-o', 'BatchMode=yes', '-o', 'StrictHostKeyChecking=yes',
                        '-o', 'ConnectTimeout=15', '-o', 'ServerAliveInterval=15',
                        '-o', 'ServerAliveCountMax=3']

    def command(self, command, marker='DS_OK', timeout=7200):
        result = subprocess.run(['ssh', '-T', *self.options, self.alias, command],
                                check=True, capture_output=True, text=True, timeout=timeout)
        lines = result.stdout.replace('\r', '').splitlines()
        # RouterOS can report a console error without a nonzero SSH status.
        if marker not in lines:
            raise RuntimeError('Router command did not return its completion marker; inspect the router')
        return lines

    def inspect(self):
        commands = [':put "DS_STATE"', ':put [/system resource get version]']
        for profile in PROFILES:
            selector = rule(profile)
            commands += [
                f':if ([:len {selector}] != 1) do={{ :error "Expected one dedicated rule" }}',
                f':if (([/ip firewall filter get {selector} chain] != "forward") or ([/ip firewall filter get {selector} action] != "drop")) do={{ :error "Unexpected dedicated rule" }}',
                f':put [/ip firewall filter get {selector} src-address-list]',
            ]
        commands += [':put [/ip firewall address-list print count-only]', ':put "DS_OK"']
        lines = self.command('; '.join(commands))
        fields = lines[lines.index('DS_STATE') + 1:lines.index('DS_OK')]
        if len(fields) != 4 or not fields[3].isdigit():
            raise ValueError('Unexpected router state response')
        version = re.fullmatch(r'7\.(\d+)(?:\.\d+)?(?: \((?:stable|long-term)\))?', fields[0])
        if not version or int(version[1]) < 16:
            raise ValueError('This helper requires a stable RouterOS v7 release with the 7.16+ command baseline')
        active = {}
        for profile, name in zip(PROFILES, fields[1:3]):
            if name not in (list_name(profile, 'A'), list_name(profile, 'B')):
                raise ValueError('Dedicated rule references an unexpected list')
            active[profile] = name
        return active, int(fields[3])

    def members(self, name):
        if not re.fullmatch(r'DS_(STANDARD|CRITICAL)_[AB]', name):
            raise ValueError('Invalid list name')
        command = (':put "DS_MEMBERS"; '
                   f':foreach item in=[/ip firewall address-list find where list="{name}"] do={{ '
                   ':if (([/ip firewall address-list get $item disabled]) or ([/ip firewall address-list get $item dynamic])) do={ :error "Unexpected disabled or dynamic entry" }; '
                   ':put [/ip firewall address-list get $item address] }; :put "DS_OK"')
        lines = self.command(command)
        values = lines[lines.index('DS_MEMBERS') + 1:lines.index('DS_OK')]
        canonical = []
        for value in values:
            ip = ipaddress.IPv4Address(value)
            if str(ip) != value:
                raise ValueError('Noncanonical loaded membership')
            canonical.append(str(ip))
        if len(set(canonical)) != len(canonical):
            raise ValueError('Duplicate loaded membership')
        return set(canonical)

    def guard(self, active):
        commands = []
        for profile in PROFILES:
            selector = rule(profile)
            commands += [f':if ([:len {selector}] != 1) do={{ :error "Dedicated rule changed" }}',
                         f':if ([/ip firewall filter get {selector} src-address-list] != "{active[profile]}") do={{ :error "Active list changed" }}']
        return commands

    def stage(self, active, staged, profiles, directory):
        commands = self.guard(active)
        for name in staged.values():
            # Refuse to clear any slot used by any IPv4 firewall table.
            for table in ('filter', 'raw', 'mangle', 'nat'):
                commands.append(f':if ([:len [/ip firewall {table} find where src-address-list="{name}" or dst-address-list="{name}"]] != 0) do={{ :error "Staging list is referenced" }}')
            commands += [f':if ([:len [/ip firewall address-list find where list="{name}" and comment!="{MARKER}"]] != 0) do={{ :error "Unmanaged staging entries" }}']
        # Check both slots before removing any stale staging entry.
        commands += [f'/ip firewall address-list remove [find where list="{name}"]' for name in staged.values()]
        commands.append(':put "DS_OK"')
        self.command('; '.join(commands))
        path = directory / 'data-shield-stage.rsc'
        with path.open('w', encoding='ascii') as handle:
            for profile in PROFILES:
                for address in profiles[profile]:
                    # Only canonical IPv4 tokens from publication() enter commands.
                    handle.write(f'/ip firewall address-list add list={staged[profile]} address={address} comment="{MARKER}"\n')
        subprocess.run(['scp', '-O', *self.options, str(path), self.alias + ':data-shield-stage.rsc'],
                       check=True, capture_output=True, text=True, timeout=7200)
        self.command(':do { /import file-name=data-shield-stage.rsc } on-error={ :error "Stage import failed" }; :put "DS_OK"')

    def activate(self, active, staged, profiles):
        commands = self.guard(active)
        for profile in PROFILES:
            commands.append(f':if ([/ip firewall address-list print count-only where list="{staged[profile]}"] != {len(profiles[profile])}) do={{ :error "Staging count changed" }}')
        switch = '; '.join(f'/ip firewall filter set {rule(p)} src-address-list="{staged[p]}"' for p in PROFILES)
        restore = '; '.join(f'/ip firewall filter set {rule(p)} src-address-list="{active[p]}"' for p in PROFILES)
        commands += [':do { ' + switch + ' } on-error={ ' + restore + '; :error "Activation failed" }',
                     ':put "DS_OK"']
        self.command('; '.join(commands))


def synchronize(router, profiles, maximum, directory, before_mutation=lambda: None):
    active, total = router.inspect()
    if all(router.members(active[p]) == set(profiles[p]) for p in PROFILES):
        return 'Both active lists already match; no import or configuration write was needed.'
    staged = {p: list_name(p, 'B' if active[p].endswith('_A') else 'A') for p in PROFILES}
    inactive = sum(len(router.members(name)) for name in staged.values())
    required = total - inactive + sum(len(profiles[p]) for p in PROFILES)
    if required > maximum:
        raise ValueError(f'Entry budget exceeded: {required} total staged and active entries need more than {maximum}')
    before_mutation()
    router.stage(active, staged, profiles, directory)
    for profile in PROFILES:
        if router.members(staged[profile]) != set(profiles[profile]):
            raise ValueError(profile + ': complete staged membership differs; activation refused')
    router.activate(active, staged, profiles)
    current, _ = router.inspect()
    if current != staged:
        raise RuntimeError('Activation could not be confirmed; inspect dedicated rules')
    for profile in PROFILES:
        if router.members(current[profile]) != set(profiles[profile]):
            raise RuntimeError('Active membership could not be confirmed; inspect the router')
    return 'Both active lists verified. Previous slots retained for rollback.'


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--apply', action='store_true', help='Explicitly allow router list updates')
    parser.add_argument('--router', help='Dedicated OpenSSH Host alias')
    parser.add_argument('--max-total-entries', type=int, help='Operator-verified capacity for all current and staging entries')
    parser.add_argument('--state-dir', type=Path, help='Private persistent state directory on the Linux controller')
    args = parser.parse_args()
    if args.apply and (not args.router or not args.state_dir or not args.max_total_entries or not 1 <= args.max_total_entries <= 2000000):
        parser.error('--apply requires --router, --state-dir and a validated --max-total-entries budget')
    try:
        generation, expected, profiles = publication()
        print('Verified publication: ' + generation['source_commit'])
        for profile in PROFILES:
            print(f'{profile.title()}: {len(profiles[profile])} records; SHA-256 {expected[profile]["sha256"]}')
        if not args.apply:
            print('Validation only. No router connection or configuration change.')
            return
        args.state_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
        with (args.state_dir / 'update.lock').open('a') as lock:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
            pending = args.state_dir / 'pending.json'
            if pending.exists():
                raise RuntimeError('An earlier apply was not confirmed. Inspect the router and running imports before clearing pending.json manually.')
            def record_pending():
                with pending.open('x') as handle:
                    json.dump({'started_at': datetime.now(timezone.utc).isoformat(), 'source_commit': generation['source_commit']}, handle)
                    handle.flush()
                    os.fsync(handle.fileno())
            # Any exception leaves a persistent stop marker, including an uncertain SSH outcome.
            with tempfile.TemporaryDirectory(prefix='routeros-', dir=args.state_dir) as temporary:
                message = synchronize(Router(args.router), profiles, args.max_total_entries, Path(temporary), record_pending)
            pending.unlink(missing_ok=True)
            print(message)
    except (ValueError, KeyError, TypeError, OSError, RuntimeError, subprocess.SubprocessError) as error:
        parser.exit(1, 'RouterOS update failed: ' + str(error) + '\n')


if __name__ == '__main__':
    main()