#!/usr/bin/env python3
# decode_ems.py - best-effort UDH/EMS object dumper
import sys
from binascii import unhexlify
from PIL import Image

def hex2bytes(h):
    h = h.strip().replace(" ", "").replace("\n","")
    if h.startswith("0x"): h = h[2:]
    return unhexlify(h)

def parse_udhl(ud):
    # ud is bytes of user data (starts with UDHL if TP-UDHI=1)
    if len(ud) == 0:
        return None
    udhl = ud[0]
    if udhl + 1 > len(ud):
        raise ValueError("UDHL longer than UD")
    udh = ud[1:1+udhl]
    rest = ud[1+udhl:]
    ies = []
    i = 0
    while i < len(udh):
        iei = udh[i]
        i += 1
        iedl = udh[i]
        i += 1
        ied = udh[i:i+iedl]
        i += iedl
        ies.append((iei, ied))
    return udhl, ies, rest


IEI_NAMES = {
    0x00: "Concatenated SMS (8-bit ref)",
    0x01: "Special SMS Message Indication",
    0x04: "Application port addressing (8-bit)",
    0x05: "Application port addressing (16-bit)",
    0x06: "SMSC Control Parameters",
    0x07: "UDH Source Indicator",
    0x08: "Concatenated SMS (16-bit ref)",
    0x0A: "Text Formatting (EMS)",
    0x0B: "Predefined Sound (EMS)",
    0x0C: "User Defined Sound (iMelody)",
    0x0D: "Predefined Animation (EMS)",
    0x0E: "Large Animation / Logo (EMS)",
    0x0F: "Small Animation (EMS)",
    0x10: "Large Picture (EMS)",
    0x11: "Small Picture (EMS)",
    0x12: "Variable Picture (EMS)",
    0x13: "User Prompt Indicator",
    0x14: "Extended Object (EMS)",
    0x15: "Reused Extended Object",
    0x18: "WVG Vector Graphic Object",
}


def decode_ie(iei, ied):
    """Decode a single IE (Information Element). Return (name, info_dict).

    For many EMS/UDH types we do a best-effort parse and otherwise return raw
    hex bytes.
    """
    info = {"raw": ied.hex()}
    name = IEI_NAMES.get(iei, f"IEI_0x{iei:02X}")
    try:
        if iei == 0x00 and len(ied) == 3:
            # concatenated short message, 8-bit ref
            info["ref8"] = ied[0]
            info["total_segments"] = ied[1]
            info["seq_number"] = ied[2]
        elif iei == 0x08 and len(ied) == 4:
            # concatenated, 16-bit ref
            info["ref16"] = (ied[0] << 8) | ied[1]
            info["total_segments"] = ied[2]
            info["seq_number"] = ied[3]
        elif iei == 0x04 and len(ied) == 2:
            info["port_dest"] = ied[0]
            info["port_orig"] = ied[1]
        elif iei == 0x05 and len(ied) == 4:
            info["port_dest"] = (ied[0] << 8) | ied[1]
            info["port_orig"] = (ied[2] << 8) | ied[3]
        elif iei == 0x06:
            # SMSC control parameters: present raw and try to break into bytes
            info["params"] = [b for b in ied]
        elif iei == 0x07:
            info["udh_source_indicator"] = ied.hex()
        elif iei == 0x01 and len(ied) >= 2:
            # Special SMS Message Indication: best-effort
            info["indication_type"] = ied[0]
            info["indication_count"] = ied[1]
            if len(ied) > 2:
                info["extra"] = ied[2:].hex()
        elif iei in (0x0B, 0x0D, 0x0E, 0x0F, 0x10, 0x11, 0x12, 0x14, 0x18):
            # EMS/UDH object types: best-effort parsing.
            # Conventionally first byte is object id or position; remaining bytes are object-specific data.
            if len(ied) >= 1:
                info["obj_first_byte"] = ied[0]
                info["obj_data"] = ied[1:].hex()
                # try to decode iMelody / ASCII-like payloads for some types
                try:
                    ascii_preview = ied[1:].decode("ascii")
                    # include preview if it looks textual
                    if all(32 <= ord(c) <= 126 or c in "\r\n\t" for c in ascii_preview[:80]):
                        info["obj_ascii_preview"] = ascii_preview[:200]
                except Exception:
                    pass
                # For picture/animation-like payloads, try the animation heuristic
                try_frames = try_decode_animation(ied[1:], out_prefix=f"udh_ie_{iei:02X}")
                if try_frames:
                    info["decoded_frames"] = try_frames
        elif iei == 0x0C:
            # user defined sound (iMelody) - position + melody bytes
            if len(ied) >= 2:
                info["position"] = ied[0]
                info["melody"] = ied[1:].hex()
                # iMelody is ASCII text (starting with "BEGIN:IMELODY" frequently)
                try:
                    txt = ied[1:].decode("ascii", errors="replace")
                    info["melody_text"] = txt
                except Exception:
                    pass
        elif iei == 0x14 and len(ied) >= 7:
            # Extended object: first bytes typically object header
            info["ext_obj_header"] = ied[:7].hex()
            info["ext_obj_data"] = ied[7:].hex()
        elif iei == 0x0A:
            # Text Formatting (EMS): usually a sequence of 4-byte formatting records
            recs = []
            if len(ied) % 4 == 0:
                for j in range(0, len(ied), 4):
                    r = ied[j:j+4]
                    recs.append({
                        "start": r[0],
                        "length": r[1],
                        "format": r[2],
                        "unused": r[3],
                    })
                info["formatting_records"] = recs
            else:
                info["formatting_raw"] = ied.hex()
        else:
            # fallback: present some breakdowns for common lengths
            if len(ied) in (1,2,3,4):
                for idx, b in enumerate(ied):
                    info[f"b{idx}"] = b
    except Exception as e:
        info["_decode_error"] = str(e)
    return name, info


def scan_object_blocks(data):
    # Heuristic: read tag (1 byte), length (1 byte), then payload
    i = 0
    objs = []
    while i+2 <= len(data):
        tag = data[i]
        length = data[i+1]
        if i+2+length > len(data):
            break
        payload = data[i+2:i+2+length]
        objs.append((i, tag, length, payload))
        i += 2 + length
    return objs

def try_decode_animation(payload, out_prefix="frame"):
    # Heuristic: assume header: frame_count, maybe width, height or similar
    # Try to find: [frame_count (1), width (1), height (1), ...frames...]
    if len(payload) < 4:
        return False
    fc = payload[0]
    w = payload[1]
    h = payload[2]
    # naive sanity checks
    if fc == 0 or w == 0 or h == 0 or fc > 200 or w > 256 or h > 256:
        return False
    # naive frame size calc: see if remaining len fits fc * ceil(w*h/8) or similar
    pixbits = w*h
    bytes_per_frame = (pixbits + 7)//8
    expected = 3 + fc * bytes_per_frame
    if len(payload) < expected:
        return False
    frames = []
    off = 3
    for f in range(fc):
        frame_bytes = payload[off:off+bytes_per_frame]
        off += bytes_per_frame
        # unpack bits to pixels (1=black,0=white)
        img = Image.new("1", (w, h))
        bits = []
        for b in frame_bytes:
            for bit in range(7, -1, -1):
                bits.append((b >> bit) & 1)
        bits = bits[:w*h]
        img.putdata([0 if bit else 255 for bit in bits])  # invert if necessary
        fname = f"{out_prefix}_{f:03d}.png"
        img.save(fname)
        frames.append(fname)
    return frames

def main():
    if len(sys.argv) < 2:
        print("Usage: decode_ems.py <user-data-hex-or-pdu-hex>")
        return
    ud_hex = sys.argv[1]
    ud = hex2bytes(ud_hex)
    print("Total UD bytes:", len(ud))
    try:
        udhl, ies, rest = parse_udhl(ud)
        print("UDHL:", udhl)
        print("IEs:")
        for iei, ied in ies:
            name, info = decode_ie(iei, ied)
            print(f"  IEI 0x{iei:02X} - {name}, len={len(ied)}")
            # pretty print info dict
            for k, v in info.items():
                print(f"    {k}: {v}")
    except Exception as e:
        print("No UDH or parsing failed:", e)
        return
    print("Remaining UD after UDH:", len(rest))
    objs = scan_object_blocks(rest)
    print("Found", len(objs), "object-like blocks (tag,len)")
    for off, tag, length, payload in objs:
        print(f" obj @ {off}: tag=0x{tag:02X} len={length}")
        # Show first bytes
        print("  payload head:", payload[:8].hex())
        # Try decode as animation
        frames = try_decode_animation(payload, out_prefix=f"obj_tag_{tag:02X}_at_{off}")
        if frames:
            print("  -> decoded animation frames:", frames)
        else:
            print("  -> not recognized as simple packed animation by heuristic")

if __name__ == "__main__":
    main()