~tuxpup/pypng-detach

pypng-detach/pypng_detach/cli.py -rw-r--r-- 1.6 KiB
35b7236bGeoff Beier initial add of attachment reading utility 4 months ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
from pathlib import Path

import click

from loguru import logger

import png
import zlib

@click.command()
@click.argument("pngfile", type=click.File('rb'))
def detach(pngfile):
    reader = png.Reader(file=pngfile)
    data_cur = 0
    data_buf = None
    compressed = False
    attachment_name = None
    decoded_attachment = None
    for chunk_type, data in reader.chunks():
        if chunk_type.decode() == "atCh":
            fn_bytes = bytearray()
            for i, b in enumerate(data):
                if b != 0:
                    fn_bytes.append(b)
                else:
                    data_cur = i+1
                    break
            attachment_name = fn_bytes.decode()
            compressed = data[data_cur] == 1
            data_buf = data[data_cur+1:]
    if attachment_name:
        logger.success(f"Found attachment named {attachment_name}. {'Not compressed.' if not compressed else 'Compressed.'}")
    else:
        logger.success(f"No attachment found.")
        return

    def inflate(buf):
        decompress = zlib.decompressobj()
        inflated = decompress.decompress(buf)
        inflated += decompress.flush()
        return inflated

    if compressed:
        decoded_attachment = inflate(data_buf)
    else:
        decoded_attachment = data_buf

    outpath = Path(attachment_name)

    if not outpath.exists():
        with open(outpath, 'wb') as out:
            out.write(decoded_attachment)
        logger.success(f"Decoded attachment written to {outpath}")
    else:
        logger.error(f"Cowardly refusing to write attachment to existing file {outpath}")