send a file through terminal qrcode

qr-send.py

#!/usr/bin/env python
import argparse
import base64
import hashlib
import json
import math
import os
import struct
import time

import segno

CAPICITY = {
    1: [17, 14, 11, 7],
    2: [32, 26, 20, 14],
    3: [53, 42, 32, 24],
    4: [78, 62, 46, 34],
    5: [106, 84, 60, 44],
    6: [134, 106, 74, 58],
    7: [154, 122, 86, 64],
    8: [192, 152, 108, 84],
    9: [230, 180, 130, 98],
    10: [271, 213, 151, 119],
    11: [321, 251, 177, 137],
    12: [367, 287, 203, 155],
    13: [425, 331, 241, 177],
    14: [458, 362, 258, 194],
    15: [520, 412, 292, 220],
    16: [586, 450, 322, 250],
    17: [644, 504, 364, 280],
    18: [718, 560, 394, 310],
    19: [792, 624, 442, 338],
    20: [858, 666, 482, 382],
    21: [929, 711, 509, 403],
    22: [1003, 779, 565, 439],
    23: [1091, 857, 611, 461],
    24: [1171, 911, 661, 511],
    25: [1273, 997, 715, 535],
    26: [1367, 1059, 751, 593],
    27: [1465, 1125, 805, 625],
    28: [1528, 1190, 868, 658],
    29: [1628, 1264, 908, 698],
    30: [1732, 1370, 982, 742],
    31: [1840, 1452, 1030, 790],
    32: [1952, 1538, 1112, 842],
    33: [2068, 1628, 1168, 898],
    34: [2188, 1722, 1228, 958],
    35: [2303, 1809, 1283, 983],
    36: [2431, 1911, 1351, 1051],
    37: [2563, 1989, 1423, 1093],
    38: [2699, 2099, 1499, 1139],
    39: [2809, 2213, 1579, 1219],
    40: [2953, 2331, 1663, 1273],
}
LEVELS = ["L", "M", "Q", "H"]
NUM_TYPE_MAP = {
    "u16": ">H",
    "u32": ">I",
    "u64": ">Q",
}
NUM_SIZE_MAP = {
    "u16": 2,
    "u32": 4,
    "u64": 8,
}


def prepend_id(number, buf, numtype):
    id_bytes = struct.pack(NUM_TYPE_MAP[numtype], number)
    return id_bytes + buf


def append_hash(buf, hash_len):
    h = hashlib.blake2b(digest_size=hash_len)
    h.update(buf)
    return buf + h.digest()


def make_hashed_stream(buf, hash_len, capicity, prefix):
    content_len = capicity - hash_len
    for i in range(len(buf) // content_len + 1):
        content = buf[i * content_len : (i + 1) * content_len]
        hashed_content = append_hash(prefix + content, hash_len)
        based64_content = base64.b64encode(hashed_content)
        yield based64_content


def make_and_display_qrcode_stream(
    gen, version, correction_level, compact, rentention, total, prefix
):
    for i, content in enumerate(gen):
        code = segno.make(
            content,
            version=version,
            error=correction_level,
        )
        time.sleep(rentention / 1000)
        print(f"{prefix} QR code {i+1}/ {total}")
        code.terminal(compact=compact, border=1)


def get_hashed_file_content_with_id(
    handler, start, capicity, hash_len, id_type, prefix
):
    handler.seek(start * capicity)
    buf = handler.read(capicity)
    with_id = prepend_id(start, buf, id_type)
    with_hash = append_hash(prefix + with_id, hash_len)
    return with_hash, buf


def get_file_md5(handler):
    md5 = hashlib.md5()
    while True:
        buf = handler.read(4096)
        if not buf:
            break
        md5.update(buf)
    return md5


def parse_range(r_str, filesize, capicity):
    start, end = r_str.split(":")
    start = int(start)
    if end == "":
        end = filesize // capicity + 1
        return start, end
    end = int(end)
    if end < start:
        raise ValueError("End must be greater than start.")
    if (end - 1) * capicity > filesize:
        raise ValueError("End must be less than filesize.")
    return start, end


def main(
    filepath,
    version,
    correction_level,
    rentention,
    compact,
    range_,
    id_type,
    hash_len,
    md5,
):
    qrcode_capicity = math.floor(
        CAPICITY[version][LEVELS.index(correction_level)] * 0.62
    )  # for base64
    filesize = os.path.getsize(filepath)
    qrcode_data_capicity = qrcode_capicity - NUM_SIZE_MAP[id_type] - hash_len - 1
    start, end = parse_range(range_, filesize, qrcode_data_capicity)
    data_qrcode_count = end - start

    metadata = {
        "qrcode_count": data_qrcode_count,
        "id_type": id_type,
        "hash_len": hash_len,
    }
    metadata_str = json.dumps(metadata)
    metadata_bytes = metadata_str.encode("utf-8")
    print(f"Data QRCode Count: {data_qrcode_count}")
    make_and_display_qrcode_stream(
        make_hashed_stream(metadata_bytes, hash_len, qrcode_capicity, b"M"),
        version,
        correction_level,
        compact,
        rentention,
        len(metadata_bytes) // (qrcode_capicity - hash_len - 1) + 1,
        "Metadata",
    )
    time.sleep(rentention / 1000)

    if md5:
        m = hashlib.md5()
    f = open(filepath, "rb")
    for i in range(start, end):
        with_id, raw = get_hashed_file_content_with_id(
            f, i, qrcode_data_capicity, hash_len, id_type, prefix=b"D"
        )
        base64_with_id = base64.b64encode(with_id)
        time.sleep(rentention / 1000)
        code = segno.make(base64_with_id, version=version, error=correction_level)
        print(f"Data QR code {i+1}/{data_qrcode_count + start}")
        code.terminal(compact=compact, border=1)
        if md5:
            m.update(raw)

    if md5:
        print(f"MD5: {m.hexdigest()}")
        make_and_display_qrcode_stream(
            make_hashed_stream(m.digest(), hash_len, qrcode_capicity, b"H"),
            version,
            correction_level,
            compact,
            rentention,
            len(m.digest()) // (qrcode_capicity - hash_len - 1) + 1,
            "MD5",
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Split file into QR codes.")
    parser.add_argument("filepath", type=str, help="Path to the file to split.")
    parser.add_argument("--version", "-v", type=int, default=6, help="QR code version.")
    parser.add_argument(
        "--correction-level",
        "-l",
        type=str,
        default="L",
        help="Error correction level.",
    )
    parser.add_argument(
        "--rentention",
        "-t",
        type=int,
        default=100,
        help="Miliseconds to wait between each QR code.",
    )
    parser.add_argument(
        "--compact",
        "-c",
        action="store_true",
        help="Compact QR code representation.",
        default=False,
    )
    parser.add_argument(
        "--range", "-r", default="0:", help="Start number of QR code.", type=str
    )
    parser.add_argument("--id-type", "-i", type=str, default="u16", help="Type of id.")
    parser.add_argument(
        "--hash-len", "-hl", type=int, default=4, help="Length of hash."
    )
    parser.add_argument("--md5", "-m", action="store_true", help="Add md5 hash.")
    args = parser.parse_args()
    main(
        args.filepath,
        args.version,
        args.correction_level,
        args.rentention,
        args.compact,
        args.range,
        args.id_type,
        args.hash_len,
        args.md5,
    )