# Copyright 2025 Eric Urban
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS AS IS AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import gzip
import zlib
import io
import struct
import sys
import tarfile
import hashlib
import subprocess as sp
import shutil
import tempfile
import os
import time
import os.path
import stat
import multiprocessing
import multiprocessing.pool
import base64

# example of making a key
# openssl genrsa -out private-key.pem 4096
# openssl rsa -in private-key.pem -pubout -out public-key.pem
# example of signing a file
# openssl dgst -sha1 -sign private-key.pem -out foo control.tar.gz

SIGTYPE = 'RSA'
OPENSSL = shutil.which('openssl')
MKSQUASHFS = shutil.which('mksquashfs')
FAST_TMP_DIR = os.environ.get('FAST_TMP_DIR', '/dev/shm')

def decompress(data):
    """Decompress a gzip compressed string in one shot.
    Return the list of decompressed sections.
    """
    while True:
        fp = io.BytesIO(data)
        if gzip._read_gzip_header(fp) is None:
            return
        # Use a zlib raw deflate compressor
        do = zlib.decompressobj(wbits=-zlib.MAX_WBITS)
        # Read all the data except the header
        decompressed = do.decompress(data[fp.tell():])
        if not do.eof or len(do.unused_data) < 8:
            raise EOFError("Compressed file ended before the end-of-stream "
                           "marker was reached")
        crc, length = struct.unpack("<II", do.unused_data[:8])
        if crc != zlib.crc32(decompressed):
            raise BadGzipFile("CRC check failed")
        if length != (len(decompressed) & 0xffffffff):
            raise BadGzipFile("Incorrect length of data produced")
        yield(decompressed)
        data = do.unused_data[8:].lstrip(b"\x00")


def convert_apk_to_nocomp(input_apk, priv_key_path, dest_path):
    # first section is signature, second is package metadata, third is package data
    with open(input_apk, 'rb') as fin:
        parts = list(decompress(fin.read()))

    # the signature won't be valid any longer, so we can ignore it
    parts[0] = None

    # compress the data section with no compression level
    parts[2] = gzip.compress(parts[2], compresslevel=0)
    # compute the sha256 hash of the data section after being compressed
    datahash = hashlib.sha256(parts[2]).hexdigest()
    # update the control section to have this new datahash

    # the control section is a tarfile, so open it up
    controlio = io.BytesIO()

    controlparts = []
    with tarfile.TarFile(name=None, mode='r', fileobj=io.BytesIO(parts[1])) as controltarin:
        for tp in controltarin:
            data = controltarin.extractfile(tp.name).read()
            controlparts.append((tp, data,))

    for i, (tp, data,) in enumerate(controlparts):
        if tp.name != '.PKGINFO':
            continue

        pkginfo = data.decode('utf-8')
        pkginfo = pkginfo.split("\n")
        for j, l in enumerate(pkginfo):
            if l.startswith('datahash = '):
                break
        else:
            raise RuntimeError("pkginfo section has no datahash to replace")
        del l

        pkginfo[j] = 'datahash = ' + datahash
        del j
        pkginfo = "\n".join(pkginfo).encode('utf-8')
        tp.size = len(pkginfo)
        tp.mtime = time.time()
        controlparts[i] = (tp, pkginfo,)
        del i
        break

    # convert controlparts back into a tar file in the same order
    controlio = io.BytesIO()
    with tarfile.TarFile(name=None, mode='w', fileobj=controlio, format=tarfile.PAX_FORMAT) as controltar:
        for tp, data in controlparts:
            controltar.addfile(tp, io.BytesIO(data))
            controloffset = controlio.tell() # save the last offset

    controlio.seek(0)
    controltar = controlio.read()
    del controlio
    # cut the last two blocks off, which are the EOF records for tar
    controltar = controltar[:controloffset]
    del controloffset

    # compress the control section with no compression level
    parts[1] = gzip.compress(controltar, compresslevel=0)
    del controltar
    controlhash = base64.b64encode(hashlib.sha1(parts[1]).digest()).decode('utf-8')

    # compute the signature of the control file
    # TODO - set tempfile to dev/shm somehow
    with tempfile.NamedTemporaryFile(dir=FAST_TMP_DIR) as tmpin:
        with tempfile.NamedTemporaryFile(dir=FAST_TMP_DIR) as tmpout:
            tmpin.write(parts[1])
            tmpin.flush()
            cmd = [OPENSSL, 'dgst', '-sha1' ,'-sign', priv_key_path, '-out', tmpout.name, tmpin.name]

            sp.check_call(cmd, executable=cmd[0])
            tmpout.seek(0)
            controlsig = tmpout.read()


    # TODO - make this configurable
    sigfilename = '.SIGN.' + SIGTYPE + '.' + os.environ['USER'] + '@none.test' + '-' + os.path.basename(priv_key_path) + '.key.pub'

    # create a tar file that is the signature
    sigio = io.BytesIO()
    with tarfile.TarFile(name=None, mode='w', fileobj=sigio, format=tarfile.GNU_FORMAT) as sigtar:
        tinfo = tarfile.TarInfo(sigfilename)
        tinfo.size = len(controlsig)
        tinfo.mtime = time.time()
        sigtar.addfile(tinfo, io.BytesIO(controlsig))
        sigoffset = sigio.tell()
    sigio.seek(0)
    sigtar = sigio.read()
    del sigio
    # cut the last two blocks off, which are the EOF records for tar
    sigtar = sigtar[:sigoffset]
    parts[0] = gzip.compress(sigtar, compresslevel=0)

    total = 0
    with open(dest_path, mode='wb') as fout:
        for part in parts:
            total += fout.write(part)
    return controlhash, total

def prepare_all_nocomp_apk(priv_key_path, alpine_arch, repos_root, dstpath):
    os.umask(0o22)
    with multiprocessing.pool.Pool() as pool:
        results = []
        for entry in os.listdir(repos_root):
            entry_stat = os.stat(os.path.join(repos_root, entry))
            if not stat.S_ISDIR(entry_stat.st_mode):
                continue

            # check for APKINDEX file
            apkindex_tar_gz_path = os.path.join(repos_root, entry, alpine_arch, 'APKINDEX.tar.gz')
            if not os.path.exists(apkindex_tar_gz_path):
                continue

            sys.stdout.write("found APKINDEX file at %s\n" % (apkindex_tar_gz_path,))

            apkindex_parts = {}
            with open(apkindex_tar_gz_path, mode='rb') as fin:
                index_tar_data = fin.read()
            index_tar_data = gzip.decompress(index_tar_data)
            with tarfile.TarFile(name=None, fileobj=io.BytesIO(index_tar_data), mode='r',) as apkindex_tar_in:
                for tp in apkindex_tar_in:
                    data = apkindex_tar_in.extractfile(tp.name).read()
                    apkindex_parts[tp.name] = (tp, data,)
            del index_tar_data


            os.makedirs(os.path.join(dstpath, entry, alpine_arch))

            newapk_index_targz_path = os.path.join(dstpath, entry, alpine_arch, 'APKINDEX.tar.gz')
            newapk_index_path = os.path.join(dstpath, entry, alpine_arch, 'APKINDEX')
            apkindex_out = open(newapk_index_path, mode='w')

            # open the APKINDEX file
            fin = io.StringIO(apkindex_parts['APKINDEX'][1].decode('utf-8'))
            accum = []
            package = None
            version = None
            checksum = None
            for line in fin:
                line = line.strip()
                if len(line) == 0:
                    data = {}
                    keys_order = []
                    for pkgdata_entry in accum:
                        k = pkgdata_entry[0]
                        v = pkgdata_entry[2:]
                        data[k] = v
                        keys_order.append(k)
                    accum = []
                    package_version = data['V']
                    package_name = data['P']
                    pkg_file_name = '%s-%s.apk' % (package_name, package_version,)
                    repodir = os.path.join(repos_root, entry, alpine_arch)

                    fullapkpath = os.path.join(repodir, pkg_file_name)
                    sys.stdout.write("found APK file at %s\n" % (fullapkpath,))
                    zerocompdstpath = os.path.join(dstpath, entry, alpine_arch, pkg_file_name)
                    sys.stdout.write("duplicating with no compression to %s\n" % (zerocompdstpath,))
                    args = (fullapkpath, priv_key_path, zerocompdstpath,)
                    results.append((data, keys_order, pool.apply_async(convert_apk_to_nocomp, args=args)))

                else:
                    accum.append(line)
            del fin

            for pkgdata, keys_order, r in results:
                controlhash, totalbytes = r.get()
                pkgdata['C'] = 'Q1' + controlhash # overwrite the hash of the package
                pkgdata['S'] = str(totalbytes) # over write the total size of the package
                for k in keys_order:
                    apkindex_out.write("%s:%s\n" % (k, pkgdata[k],))
                apkindex_out.write("\n") # blank linke to end the package definition
            apkindex_out.close()

            with tarfile.TarFile.gzopen(newapk_index_targz_path, compresslevel=0, mode='w') as tarout:
                with open(newapk_index_path, mode='rb') as fin:
                    tp = apkindex_parts['APKINDEX'][0]
                    data = fin.read()
                    tp.size = len(data)
                    tp.mtime = time.time()
                    tarout.addfile(tp, io.BytesIO(data))
                tp, data = apkindex_parts['DESCRIPTION']
                tarout.addfile(tp, io.BytesIO(data))
            os.unlink(newapk_index_path)




def build_squashfs(priv_key_path, alpine_arch, repos_root, dstpath):
    prepare_all_nocomp_apk(priv_key_path, alpine_arch, repos_root, dstpath)

    cmd = [MKSQUASHFS, dstpath, 'apk_' + alpine_arch + '.squashfs', '-no-xattrs', '-comp', 'xz', '-Xdict-size', '1M', '-b' , '1M', '-not-reproducible']
    cmd.extend(['-all-time', str(int(time.time()))])
    cmd.extend(['-force-uid', '1'])
    cmd.extend(['-force-gid', '1'])
    cmd.extend(['-root-mode', '555'])
    sp.check_call(cmd, stdin=sp.DEVNULL, executable=cmd[0])


mode = sys.argv[1]

def suffix_apk(input_apk):
    input_apk = os.path.basename(input_apk)
    head, tail = input_apk.rsplit('.', 1)
    return head + '-nocomp.' + tail

if mode == 'nocomp':
    priv_key_path = sys.argv[2]
    for input_apk in sys.argv[3:]:
        destpath = suffix_apk(input_apk)
        convert_apk_to_nocomp(input_apk, priv_key_path, destpath)
elif mode == 'squashfs':
    priv_key_path = sys.argv[2]
    alpine_arch = sys.argv[3]
    repos_root = sys.argv[4]
    with tempfile.TemporaryDirectory(dir=os.environ.get('TMPDIR')) as tmpdir:
        build_squashfs(priv_key_path, alpine_arch, repos_root, tmpdir)
else:
    sys.stderr.write("unknown mode: ")
    sys.stderr.write(mode)
    sys.stderr.write("\n")