import zipfile, hashlib, base64, struct, os, datetime
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives.serialization import pkcs7
from cryptography import x509
from cryptography.x509.oid import NameOID

CREATED_BY = "Claude Code APK repacker"

def wrap(line: bytes) -> bytes:
    """JAR manifest continuation wrapping: <=72 bytes per line incl. CRLF."""
    out, first = b"", True
    while line:
        n = 70 if first else 69
        chunk, line = line[:n], line[n:]
        out += (b"" if first else b" ") + chunk + b"\r\n"
        first = False
    return out

def section(name: str, digest_b64: str) -> bytes:
    return wrap(b"Name: " + name.encode()) + wrap(b"SHA-256-Digest: " + digest_b64.encode()) + b"\r\n"

def b64(b): return base64.b64encode(b).decode()

def make_key():
    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    now = datetime.datetime(2026, 1, 1)
    nm = x509.Name([
        x509.NameAttribute(NameOID.COMMON_NAME, "MCPE Compat Resign"),
        x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, "Local rebuild"),
    ])
    cert = (x509.CertificateBuilder()
            .subject_name(nm).issuer_name(nm)
            .public_key(key.public_key())
            .serial_number(x509.random_serial_number())
            .not_valid_before(now)
            .not_valid_after(datetime.datetime(2056, 1, 1))
            .sign(key, hashes.SHA256()))
    return key, cert

def build_v1(src_apk, patched_manifest, out_apk, key, cert):
    zin = zipfile.ZipFile(src_apk)
    entries = []
    for i in zin.infolist():
        if i.filename.startswith("META-INF/"):
            continue
        data = open(patched_manifest, 'rb').read() if i.filename == "AndroidManifest.xml" else zin.read(i)
        entries.append((i, data))

    mf = b"Manifest-Version: 1.0\r\nCreated-By: " + CREATED_BY.encode() + b"\r\n\r\n"
    sf_sections = b""
    for i, data in entries:
        sec = section(i.filename, b64(hashlib.sha256(data).digest()))
        mf += sec
        sf_sections += section(i.filename, b64(hashlib.sha256(sec).digest()))

    sf = (b"Signature-Version: 1.0\r\n"
          b"Created-By: " + CREATED_BY.encode() + b"\r\n"
          b"SHA-256-Digest-Manifest: " + b64(hashlib.sha256(mf).digest()).encode() + b"\r\n\r\n"
          + sf_sections)

    p7 = (pkcs7.PKCS7SignatureBuilder()
          .set_data(sf)
          .add_signer(cert, key, hashes.SHA256())
          .sign(serialization.Encoding.DER,
                [pkcs7.PKCS7Options.DetachedSignature, pkcs7.PKCS7Options.NoAttributes]))

    zout = zipfile.ZipFile(out_apk, 'w')
    for i, data in entries:
        zi = zipfile.ZipInfo(i.filename, date_time=i.date_time)
        zi.compress_type = i.compress_type
        zi.external_attr = i.external_attr
        zout.writestr(zi, data)
    for nm, blob in (("META-INF/MANIFEST.MF", mf), ("META-INF/CERT.SF", sf), ("META-INF/CERT.RSA", p7)):
        zi = zipfile.ZipInfo(nm, date_time=(2026, 1, 1, 0, 0, 0))
        zi.compress_type = zipfile.ZIP_DEFLATED
        zout.writestr(zi, blob)
    zout.close()
    return len(entries)

CHUNK = 1048576
def chunk_digests(blobs):
    digs = []
    for b in blobs:
        for o in range(0, len(b), CHUNK):
            c = b[o:o+CHUNK]
            h = hashlib.sha256()
            h.update(b"\xa5" + struct.pack('<I', len(c)) + c)
            digs.append(h.digest())
    top = hashlib.sha256()
    top.update(b"\x5a" + struct.pack('<I', len(digs)) + b"".join(digs))
    return top.digest()

def lp(b): return struct.pack('<I', len(b)) + b
def seq(items): return lp(b"".join(lp(x) for x in items))

def sign_v2(path, key, cert):
    d = bytearray(open(path, 'rb').read())
    # locate EOCD
    eocd_off = d.rfind(b"PK\x05\x06")
    if eocd_off < 0: raise SystemExit("no EOCD")
    cd_off = struct.unpack_from('<I', d, eocd_off + 16)[0]
    before_cd = bytes(d[:cd_off])
    cd = bytes(d[cd_off:eocd_off])
    eocd = bytes(d[eocd_off:])

    ALGO = 0x0103  # RSASSA-PKCS1-v1_5 with SHA2-256
    digest = chunk_digests([before_cd, cd, eocd])
    digests = seq([struct.pack('<I', ALGO) + lp(digest)])
    certs = seq([cert.public_bytes(serialization.Encoding.DER)])
    attrs = lp(b"")
    signed_data = digests + certs + attrs

    sig = key.sign(signed_data, padding.PKCS1v15(), hashes.SHA256())
    signatures = seq([struct.pack('<I', ALGO) + lp(sig)])
    pub = cert.public_key().public_bytes(serialization.Encoding.DER,
                                         serialization.PublicFormat.SubjectPublicKeyInfo)
    signer = lp(signed_data) + signatures + lp(pub)
    v2_value = seq([signer])

    pair = struct.pack('<Q', len(v2_value) + 4) + struct.pack('<I', 0x7109871a) + v2_value
    body = pair
    # pad so the whole signing block is a multiple of 4096 (recommended, not required)
    blk_len = len(body) + 8 + 16  # trailing size + magic
    pad = (-(blk_len + 8)) % 4096
    if pad:
        if pad < 12: pad += 4096
        body += struct.pack('<Q', pad - 8) + struct.pack('<I', 0x42726577) + b"\x00" * (pad - 12)
    size = len(body) + 8 + 16
    block = struct.pack('<Q', size) + body + struct.pack('<Q', size) + b"APK Sig Block 42"

    new_eocd = bytearray(eocd)
    struct.pack_into('<I', new_eocd, 16, cd_off + len(block))
    open(path, 'wb').write(before_cd + block + cd + bytes(new_eocd))
    return len(block)

if __name__ == '__main__':
    key, cert = make_key()
    n = build_v1('/root/pe/PE-a0.2.0-1.apk', 'AndroidManifest.patched.xml', 'PE-0.2.0-compat.apk', key, cert)
    print(f"v1: repacked {n} entries, JAR-signed (SHA-256)")
    bl = sign_v2('PE-0.2.0-compat.apk', key, cert)
    print(f"v2: APK Signing Block inserted, {bl} bytes")
    with open('signing-key.pem', 'wb') as f:
        f.write(key.private_bytes(serialization.Encoding.PEM,
                                  serialization.PrivateFormat.PKCS8,
                                  serialization.NoEncryption()))
    print("key saved to signing-key.pem")