#!/usr/bin/env python3
"""Vendor a pinned HardeningKitty GitHub snapshot on the PACKAGING workstation.

Endpoints never download executable code from GitHub. All selected upstream files
are subsequently included and signed with the WAPT package.
"""
import argparse
import hashlib
import io
import json
import os
from pathlib import Path
import shutil
import sys
import tempfile
import urllib.request
import zipfile

UPSTREAM_REPO = '0x6d69636b/windows_hardening'
UPSTREAM_COMMIT = 'acf9ecee3d044a9337df7aad80e730cbd746e8de'
UPSTREAM_API = 'https://api.github.com/repos/%s/zipball/%s' % (UPSTREAM_REPO, UPSTREAM_COMMIT)
MAX_ARCHIVE_BYTES = 30 * 1024 * 1024
MAX_FILE_BYTES = 3 * 1024 * 1024
MAX_TOTAL_EXTRACTED_BYTES = 25 * 1024 * 1024
REQUIRED = {
    'HardeningKitty.psm1',
    'HardeningKitty.psd1',
    'LICENSE',
    'lists/hardeningkitty_lists_manifest.psd1',
    'lists/hardeningkitty_lists_manifest.psd1.p7s',
    'lists/finding_list_0x6d69636b_machine.csv',
    'lists/finding_list_msft_security_baseline_windows_11_24h2_machine.csv',
    'lists/finding_list_msft_security_baseline_windows_11_25h2_machine.csv',
    'lists/finding_list_msft_security_baseline_windows_10_22h2_machine.csv',
}


def sha256(data):
    return hashlib.sha256(data).hexdigest()


def _download_snapshot():
    req = urllib.request.Request(UPSTREAM_API, headers={'User-Agent': 'comitari-wapt-hardeningkitty-vendor',
                                                       'Accept': 'application/vnd.github+json'})
    with urllib.request.urlopen(req, timeout=90) as response:
        if response.url.split('/')[0] != 'https:':
            raise RuntimeError('Refused non-HTTPS archive download.')
        data = response.read(MAX_ARCHIVE_BYTES + 1)
        if len(data) > MAX_ARCHIVE_BYTES:
            raise RuntimeError('GitHub archive exceeds size limit.')
        return data


def _wanted(relative):
    if relative in ('HardeningKitty.psm1', 'HardeningKitty.psd1', 'LICENSE'):
        return True
    if relative.startswith('lists/') and relative.count('/') == 1:
        return (relative.endswith('.csv') or relative in
                ('lists/hardeningkitty_lists_manifest.psd1',
                 'lists/hardeningkitty_lists_manifest.psd1.p7s'))
    return False


def vendor_from_archive(archive_bytes, destination):
    root = Path(destination)
    root.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix='vendor-tmp-', dir=str(root.parent)) as tmp:
        files = {}
        uncompressed = 0
        with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive:
            for entry in archive.infolist():
                if entry.is_dir():
                    continue
                segments = entry.filename.split('/')
                if len(segments) < 2:
                    continue
                if UPSTREAM_COMMIT[:7].lower() not in segments[0].lower():
                    raise RuntimeError('Unexpected GitHub archive revision: %s' % segments[0])
                relative = '/'.join(segments[1:])
                if not _wanted(relative):
                    continue
                if (relative.startswith('/') or '..' in segments or '\\' in relative or
                    not all(p and p not in ('.', '..') for p in segments[1:])):
                    raise RuntimeError('Unsafe archive path: %s' % entry.filename)
                if entry.file_size > MAX_FILE_BYTES or uncompressed + entry.file_size > MAX_TOTAL_EXTRACTED_BYTES:
                    raise RuntimeError('Unreasonable uncompressed payload size.')
                if relative in files:
                    raise RuntimeError('Duplicate upstream path: %s' % relative)
                with archive.open(entry) as src:
                    data = src.read(MAX_FILE_BYTES + 1)
                if len(data) != entry.file_size:
                    raise RuntimeError('Corrupted upstream file: %s' % relative)
                uncompressed += len(data)
                target = Path(tmp, *relative.split('/'))
                target.parent.mkdir(parents=True, exist_ok=True)
                target.write_bytes(data)
                files[relative] = sha256(data)
        missing = REQUIRED.difference(files)
        if missing:
            raise RuntimeError('Incomplete upstream archive; missing: %s' % ', '.join(sorted(missing)))
        list_count = sum(name.startswith('lists/') and name.endswith('.csv') for name in files)
        if list_count < 50:
            raise RuntimeError('Insufficient reference lists: %d' % list_count)
        manifest = {
            'upstream_repository': UPSTREAM_REPO,
            'upstream_commit': UPSTREAM_COMMIT,
            'source_archive_sha256': sha256(archive_bytes),
            'source': 'pinned GitHub archive (not a floating branch)',
            'files': dict(sorted(files.items())),
        }
        Path(tmp, 'vendor-manifest.json').write_text(json.dumps(manifest, indent=2, sort_keys=True) + '\n', encoding='utf-8')
        # Preserve a previous verified vendor folder when the download fails.
        if root.exists():
            shutil.rmtree(root)
        shutil.move(tmp, root)
    return manifest


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--dest', default='vendor', help='Directory to populate in the WAPT package sources')
    parser.add_argument('--archive', help='Use a locally downloaded GitHub archive for offline preparation')
    parser.add_argument('--archive-sha256', help='Optional expected hash of the local ZIP for offline verification')
    args = parser.parse_args(argv)
    if args.archive:
        data = Path(args.archive).read_bytes()
        if len(data) > MAX_ARCHIVE_BYTES:
            raise RuntimeError('Offline archive exceeds size limit.')
        if args.archive_sha256 and sha256(data).lower() != args.archive_sha256.lower():
            raise RuntimeError('Offline archive checksum mismatch.')
    else:
        print('Downloading pinned HardeningKitty GitHub archive on packaging workstation...')
        data = _download_snapshot()
    result = vendor_from_archive(data, args.dest)
    print('Vendored %d files, commit %s, archive SHA256 %s' %
          (len(result['files']), result['upstream_commit'], result['source_archive_sha256']))
    print('Now review, sign and build the WAPT package; endpoints need no Internet.')


if __name__ == '__main__':
    try:
        main()
    except Exception as exc:
        print('ERROR: %s' % exc, file=sys.stderr)
        sys.exit(1)