#!/usr/bin/env python3
"""Cache mainland domain/IP rule sets locally. Does not restart any service."""
import argparse
import os
import tempfile
import urllib.request
from pathlib import Path

parser = argparse.ArgumentParser()
parser.add_argument('--output', default='data/mihomo/rules')
args = parser.parse_args()
directory = Path(args.output)
directory.mkdir(parents=True, exist_ok=True)
files = {'cn-domain.mrs': 'geosite/cn.mrs', 'cn-ip.mrs': 'geoip/cn.mrs'}
downloads = {}
for name, remote in files.items():
    request = urllib.request.Request('https://raw.githubusercontent.com/MetaCubeX/meta-rules-dat/meta/geo/' + remote, headers={'User-Agent': 'ClashWeb-build'})
    with urllib.request.urlopen(request, timeout=60) as response:
        data = response.read(32 * 1024 * 1024 + 1)
    # MRS resources use a Zstandard frame; core validation checks the decoded payload.
    if len(data) > 32 * 1024 * 1024 or data[:4] != b'\x28\xb5\x2f\xfd':
        raise SystemExit(f'{name}: invalid MRS resource or size exceeds limit')
    downloads[name] = data
for name, data in downloads.items():
    fd, temporary = tempfile.mkstemp(dir=directory, prefix='.rules-')
    try:
        with os.fdopen(fd, 'wb') as stream:
            stream.write(data)
        os.chmod(temporary, 0o640)
        os.replace(temporary, directory / name)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)
    print(f'Cached {name}: {len(data)} bytes')