multidecoder.py

from __future__ import annotations

import base64
import html
import sys
import urllib.parse
from typing import Callable, Iterable, NewType

DECODERS: dict[str, Callable[[bytes], bytes]] = {}


def add_decoder(name: str):
    def decorator(fn):
        DECODERS[name] = fn
        return fn

    return decorator


@add_decoder("base64")
def decode_base64(text: bytes) -> bytes:
    return base64.decodebytes(text)


@add_decoder("base64_urlsafe")
def decode_base64_urlsafe(text: bytes) -> bytes:
    return base64.urlsafe_b64decode(text)


@add_decoder("hex")
def decode_hex(text: bytes) -> bytes:
    return bytes.fromhex(text.decode())


@add_decoder("html")
def decode_html(text: bytes) -> bytes:
    return html.unescape(text)


@add_decoder("url")
def decode_url(text: bytes) -> bytes:
    return urllib.parse.unquote(text.decode()).encode()


Result = NewType("Result", tuple[str, bytes, list["Result"]])


def multidecode(text: bytes, max_depth=10) -> Iterable[Result]:
    if max_depth <= 0:
        return  # Max depth reached

    assert isinstance(text, bytes)
    for name, decoder in DECODERS.items():
        try:
            decoded = decoder(text)
        except Exception:
            continue  # Nothing was decoded
        assert isinstance(decoded, bytes)
        if text == decoded:
            continue  # Nothing changed

        yield Result((name, decoded, multidecode(decoded, max_depth=max_depth - 1)))


def display_result(results: Iterable[Result], depth=0):
    indent = " " * (4 * depth)
    for (
        name,
        decoded,
        subitems,
    ) in results:
        print(f"{indent}{name} -> {decoded}")
        display_result(subitems, depth + 1)


def main():
    for text in sys.argv[1:]:
        print(text)
        text = text.encode()
        results = multidecode(text, max_depth=10)
        display_result(results, depth=1)
        print()


if __name__ == "__main__":
    main()