strings(1) with .tbl file

tblstrings.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""strings(1) with .tbl file

Example:

    $ tblstrings --tbl game.tbl game.nes
    $ tblstrings --tbl game.tbl -tx game.nes # show offsets

For .tbl files, see: http://www.fceux.com/web/help/fceux.html?HexEditor.html
"""

import argparse
import io
import re
import sys

def chomp(s):
    return s.rstrip("\r\n")

def iter_bytes(in_, bufsiz=io.DEFAULT_BUFFER_SIZE):
    b = in_.read(bufsiz)
    while b:
        yield from b
        b = in_.read(bufsiz)

BYTES_DEF = 4

TBL_DEF = dict((i,chr(i)) for i in range(0x20,0x7F)) # ASCII printable

RE_TBL_LINE = re.compile(r"\A([0-9a-fA-F]{1,2})=(.*)\Z")

def strings(in_, name, opts):
    for off, s in iter_strings(in_, opts.tbl, opts.bytes_):
        if opts.print_file_name:
            print(f"{name}:\t", end="")
        if opts.radix is not None:
            print(f"{str_radix(off,opts.radix)}\t", end="")
        print(s)

def iter_strings(in_, tbl, len_min):
    buf = bytearray()
    for i, b in enumerate(iter_bytes(in_)):
        if b in tbl:
            buf.append(b)
        else:
            if len(buf) >= len_min:
                yield i-len(buf), "".join(tbl[bb] for bb in buf)
            buf.clear()

    if len(buf) >= len_min:
        yield i-len(buf), "".join(tbl[b] for b in buf)

def str_radix(n, radix):
    if radix == 8:
        return f"{n:o}"
    elif radix == 10:
        return f"{n}"
    elif radix == 16:
        return f"{n:x}"
    else:
        assert False

def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("-f", "--print-file-name", action="store_true")
    ap.add_argument("-n", "--bytes", dest="bytes_", type=arg_bytes, default=BYTES_DEF)
    ap.add_argument("-t", "--radix", type=arg_radix)
    ap.add_argument("-x", "--tbl", type=arg_tbl, default=TBL_DEF)
    ap.add_argument("files", nargs=argparse.REMAINDER)
    return ap.parse_args()

def arg_bytes(s):
    try:
        res = int(s)
    except ValueError as e:
        raise argparse.ArgumentTypeError(f"{e}")
    if res <= 0:
        raise argparse.ArgumentTypeError("bytes must be positive")
    return res

def arg_radix(s):
    MAP = dict(o=8, d=10, x=16)
    res = MAP.get(s)
    if res is None:
        raise argparse.ArgumentTypeError(f"invalid radix ({MAP})")
    return res

def arg_tbl(s):
    try:
        in_ = open(s, "r")
    except OSError as e:
        raise argparse.ArgumentTypeError(f"{e}")

    res = {}
    with in_:
        for line in map(chomp, in_):
            if not line.strip(): continue
            m = RE_TBL_LINE.fullmatch(line)
            if not m:
                raise argparse.ArgumentTypeError(f"invalid tbl line: {line}")
            k = int(m.group(1), base=16)
            v = m.group(2)
            res[k] = v

    return res

def open_file(s):
    if s == "-":
        return sys.stdin.buffer, "{standard input}"
    else:
        return open(s, "rb"), s

def main():
    args = parse_args()

    if not args.files:
        args.files.append("-")

    for s in args.files:
        in_, name = open_file(s)
        with in_:
            strings(in_, name, args)

if __name__ == "__main__": main()