import struct
import zlib
import os
import sys

# ── Config ─────────────────────────────────────────────────────────────────────
FILENAME   = "LGBackup_FILENAME.lbf"
OUTPUT_DIR = "extracted"
# ───────────────────────────────────────────────────────────────────────────────

LOCAL_SIG       = b'PK\x03\x04'
DATA_DESC_SIG   = b'PK\x07\x08'

# Local file header layout (after the 4-byte signature):
#   version_needed  2B  (H)
#   flags           2B  (H)
#   compression     2B  (H)
#   mod_time        2B  (H)
#   mod_date        2B  (H)
#   crc32           4B  (I)
#   compressed_sz   4B  (I)
#   uncompressed_sz 4B  (I)
#   fname_len       2B  (H)
#   extra_len       2B  (H)
HEADER_FMT  = '<HHHHHIIIHH'
HEADER_SIZE = struct.calcsize(HEADER_FMT)   # 26 bytes

COMPRESSION_STORED  = 0
COMPRESSION_DEFLATE = 8


def scan_to_signature(f, sig):
    """Advance f until `sig` is found; return file offset of first byte of sig,
    or -1 if EOF is reached first."""
    buf = b''
    n   = len(sig)
    while True:
        byte = f.read(1)
        if not byte:
            return -1
        buf = (buf + byte)[-n:]
        if buf == sig:
            return f.tell() - n


def read_local_header(f):
    """Parse a local file header that starts at the current position (after sig)."""
    raw = f.read(HEADER_SIZE)
    if len(raw) < HEADER_SIZE:
        return None
    fields = struct.unpack(HEADER_FMT, raw)
    (version_needed, flags, compression, mod_time, mod_date,
     crc32, compressed_sz, uncompressed_sz, fname_len, extra_len) = fields

    fname = f.read(fname_len).decode('utf-8', errors='replace')
    f.read(extra_len)   # skip extra field

    return {
        'flags':           flags,
        'compression':     compression,
        'crc32':           crc32,
        'compressed_sz':   compressed_sz,
        'uncompressed_sz': uncompressed_sz,
        'fname':           fname,
        'has_data_desc':   bool(flags & 0x0008),
    }


def read_data_with_descriptor(f):
    """When sizes are 0 (data descriptor flag), slurp bytes until we find
    PK\\x07\\x08 or the next PK\\x03\\x04, then back up past the descriptor."""
    buf = bytearray()
    while True:
        byte = f.read(1)
        if not byte:
            # Reached EOF — return whatever we have
            return bytes(buf)
        buf += byte
        tail = bytes(buf[-4:])

        if tail == DATA_DESC_SIG:
            # Consume the 12-byte descriptor body (crc32 + comp_sz + uncomp_sz)
            f.read(12)
            return bytes(buf[:-4])

        if tail == LOCAL_SIG:
            # Stumbled into the next entry — rewind
            f.seek(-4, 1)
            return bytes(buf[:-4])


def safe_output_path(base_dir, fname):
    """Resolve an Android absolute path safely under base_dir."""
    # Strip leading slashes so os.path.join doesn't discard base_dir
    fname = fname.lstrip('/')
    # Replace internal@/external@ storage prefixes with readable folder names
    fname = fname.replace('internal@/', 'internal/').replace('external@/', 'external/')
    target = os.path.realpath(os.path.join(base_dir, fname))
    if not target.startswith(os.path.realpath(base_dir)):
        raise ValueError(f"Path traversal attempt blocked: {fname!r}")
    return target


def extract(lbf_path, out_dir):
    os.makedirs(out_dir, exist_ok=True)
    extracted = skipped = 0

    with open(lbf_path, 'rb') as f:
        while True:
            pos = scan_to_signature(f, LOCAL_SIG)
            if pos < 0:
                break   # No more entries

            f.seek(pos + 4)   # skip past the 4-byte signature we just found
            header = read_local_header(f)
            if header is None:
                continue

            fname       = header['fname']
            compression = header['compression']

            # ── Directory entry ───────────────────────────────────────────────
            if fname.endswith('/'):
                os.makedirs(safe_output_path(out_dir, fname), exist_ok=True)
                continue

            # ── Read compressed payload ───────────────────────────────────────
            if header['has_data_desc'] and header['compressed_sz'] == 0:
                compressed_data = read_data_with_descriptor(f)
            else:
                compressed_data = f.read(header['compressed_sz'])

            # ── Decompress ────────────────────────────────────────────────────
            try:
                if compression == COMPRESSION_STORED:
                    raw = compressed_data
                elif compression == COMPRESSION_DEFLATE:
                    raw = zlib.decompress(compressed_data, -15)   # raw deflate
                else:
                    print(f"  [SKIP] {fname}  (unsupported compression={compression})")
                    skipped += 1
                    continue
            except zlib.error as e:
                print(f"  [ERR]  {fname}  ({e})")
                skipped += 1
                continue

            # ── Write ─────────────────────────────────────────────────────────
            try:
                out_path = safe_output_path(out_dir, fname)
                os.makedirs(os.path.dirname(out_path), exist_ok=True)
                with open(out_path, 'wb') as out:
                    out.write(raw)
                extracted += 1
                print(f"  [OK]   {fname}  ({len(raw):,} bytes)")
            except (OSError, ValueError) as e:
                print(f"  [ERR]  {fname}  ({e})")
                skipped += 1

    print(f"\n{'─'*50}")
    print(f"Extracted : {extracted}")
    print(f"Skipped   : {skipped}")
    print(f"Output    : {os.path.abspath(out_dir)}/")


if __name__ == '__main__':
    lbf = sys.argv[1] if len(sys.argv) > 1 else FILENAME
    if not os.path.exists(lbf):
        sys.exit(f"File not found: {lbf!r}")
    extract(lbf, OUTPUT_DIR)