"""Export public APK signing information and verify continuity. No private keys."""
import argparse
import hashlib
import json
import struct
import zipfile
from pathlib import Path

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.serialization import pkcs7
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography import x509


def length_prefixed(data, offset=0):
    size = struct.unpack_from('<I', data, offset)[0]
    end = offset + 4 + size
    if end > len(data):
        raise ValueError('Malformed APK signing record.')
    return data[offset + 4:end], end


def modern_certificate(apk):
    # APK v2/v3 certificates live before the ZIP central directory, not META-INF.
    with open(apk, 'rb') as stream:
        stream.seek(0, 2)
        size = stream.tell()
        stream.seek(max(0, size - 65557))
        tail = stream.read()
        eocd = tail.rfind(b'PK\x05\x06')
        if eocd < 0:
            raise ValueError('Missing ZIP directory.')
        central = struct.unpack_from('<I', tail, eocd + 16)[0]
        stream.seek(central - 24)
        footer = stream.read(24)
        if footer[8:] != b'APK Sig Block 42':
            return None
        block_size = struct.unpack_from('<Q', footer)[0]
        if block_size < 24 or block_size + 8 > central:
            raise ValueError('Malformed APK signing block.')
        stream.seek(central - block_size - 8)
        block = stream.read(block_size + 8)
    if struct.unpack_from('<Q', block)[0] != block_size:
        raise ValueError('APK signing block sizes differ.')
    offset = 8
    while offset < len(block) - 24:
        pair_size = struct.unpack_from('<Q', block, offset)[0]
        offset += 8
        pair = block[offset:offset + pair_size]
        if pair_size < 4 or len(pair) != pair_size:
            raise ValueError('Malformed APK signer pair.')
        scheme = struct.unpack_from('<I', pair)[0]
        if scheme in (0x7109871A, 0xF05368C0, 0x1B93AD61):
            signers, _ = length_prefixed(pair[4:])
            signer, signer_end = length_prefixed(signers)
            if signer_end != len(signers):
                raise ValueError('Multiple APK signers are not supported by this report.')
            signed_data, _ = length_prefixed(signer)
            _, digest_end = length_prefixed(signed_data)
            certificates, _ = length_prefixed(signed_data, digest_end)
            cert_der, _ = length_prefixed(certificates)
            return x509.load_der_x509_certificate(cert_der)
        offset += pair_size
    return None


def certificate(apk):
    modern = modern_certificate(apk)
    if modern is not None:
        return modern
    with zipfile.ZipFile(apk) as archive:
        blocks = [name for name in archive.namelist()
                  if name.upper().startswith('META-INF/')
                  and name.upper().endswith(('.RSA', '.EC', '.DSA'))]
        if len(blocks) != 1:
            raise ValueError('Expected one v1 signing certificate block; verify APK with apksigner.')
        certificates = pkcs7.load_der_pkcs7_certificates(archive.read(blocks[0]))
        if len(certificates) != 1:
            raise ValueError('Expected one signing certificate, not a certificate chain.')
        return certificates[0]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--previous', required=True)
    parser.add_argument('--current', required=True)
    parser.add_argument('--output-dir', required=True)
    args = parser.parse_args()
    previous = certificate(args.previous)
    current = certificate(args.current)
    old_der = previous.public_bytes(serialization.Encoding.DER)
    der = current.public_bytes(serialization.Encoding.DER)
    public = current.public_key().public_bytes(
        serialization.Encoding.DER, serialization.PublicFormat.SubjectPublicKeyInfo)
    if old_der != der:
        raise ValueError('Signing certificate changed. Existing signing material must be preserved.')
    output = Path(args.output_dir)
    output.mkdir(parents=True, exist_ok=True)
    details = {
        'appName': 'TPSchedule',
        'applicationId': 'com.tpschedule.app',
        'certificateUnchanged': True,
        'subject': current.subject.rfc4514_string(),
        'certificateSHA256': hashlib.sha256(der).hexdigest().upper(),
        'certificateSHA1': hashlib.sha1(der).hexdigest().upper(),
        'certificateMD5': hashlib.md5(der).hexdigest().upper(),
        'publicKeySHA256': hashlib.sha256(public).hexdigest().upper(),
        'publicKeyMD5': hashlib.md5(public).hexdigest().upper(),
        'publicKeySPKIHex': public.hex().upper(),
        'publicKeyBits': current.public_key().key_size,
        'certificateNote': 'Existing Android Debug certificate preserved; store/filing acceptance must be checked.',
    }
    if isinstance(current.public_key(), rsa.RSAPublicKey):
        numbers = current.public_key().public_numbers()
        details['publicKeyModulusDecimal'] = str(numbers.n)
        details['publicKeyExponent'] = numbers.e
    (output / 'TPSchedule-signing-info.json').write_text(
        json.dumps(details, ensure_ascii=False, indent=2), encoding='utf-8')
    (output / 'TPSchedule-signing-certificate.pem').write_bytes(
        current.public_bytes(serialization.Encoding.PEM))
    (output / 'TPSchedule-public-key.pem').write_bytes(
        current.public_key().public_bytes(serialization.Encoding.PEM,
                                         serialization.PublicFormat.SubjectPublicKeyInfo))
    print(json.dumps({'certificateUnchanged': True,
                      'certificateSHA256': details['certificateSHA256'],
                      'publicFiles': str(output)}, ensure_ascii=True))


if __name__ == '__main__':
    main()