NES Game Genie encoder/decoder

nesgenie.py

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

"""NES Game Genie encoder/decoder.

Usage:
  nesgenie enc <addr> <value> [cmp]
  nesgenie dec <code>
"""

import argparse


#---------------------------------------------------------------------
# util
#---------------------------------------------------------------------

def to_bits(value, len_):
    assert len_ >= 1
    assert value < (1<<len_)
    return tuple(1 if (value&(1<<i)) else 0 for i in reversed(range(len_)))

def from_bits(bits):
    len_ = len(bits)
    value = 0
    for i, bit in enumerate(bits):
        value |= bit << (len_-1-i)
    return value

def seq_get(seq, i, default):
    if i is None: return default
    return seq[i]


#---------------------------------------------------------------------
# genie
#---------------------------------------------------------------------

# http://nesdev.com/nesgg.txt

GENIE_MAP_6 = (
    # addr
    None, 13, 14, 15,
      16, 21, 22, 23,
       4,  9, 10, 11,
      12, 17, 18, 19,
    # value
     0, 5, 6, 7,
    20, 1, 2, 3
)

GENIE_MAP_8 = (
    # addr
    None, 13, 14, 15,
      16, 21, 22, 23,
       4,  9, 10, 11,
      12, 17, 18, 19,
    # value
     0, 5, 6, 7,
    28, 1, 2, 3,
    # cmp
    24, 29, 30, 31,
    20, 25, 26, 27
)

GENIE_CHARS = "APZLGITYEOXUKSVN"

def genie_chr(n):
    return GENIE_CHARS[n]

def genie_ord(c):
    return GENIE_CHARS.index(c)

def genie_shuffle(bits, map_):
    return tuple(seq_get(bits, map_.get(i), 0) for i in range(len(bits)))

def genie_encode_6(addr, value):
    map_rev = dict((v,k) for k,v in enumerate(GENIE_MAP_6))

    plain  = to_bits(addr, 16) + to_bits(value, 8)
    cipher = genie_shuffle(plain, map_rev)

    code = "".join(genie_chr(from_bits(cipher[i:i+4])) for i in range(0, 24, 4))
    print(code)

def genie_encode_8(addr, value, cmp_):
    map_rev = dict((v,k) for k,v in enumerate(GENIE_MAP_8))

    plain  = to_bits(addr, 16) + to_bits(value, 8) + to_bits(cmp_, 8)
    cipher = genie_shuffle(plain, map_rev)

    code = "".join(genie_chr(from_bits(cipher[i:i+4])) for i in range(0, 32, 4))
    print(code)

def genie_encode(args):
    addr  = args.addr
    value = args.value
    cmp_  = args.cmp_

    if cmp_ is None:
        genie_encode_6(addr, value)
    else:
        genie_encode_8(addr, value, cmp_)

def genie_decode_6(code):
    map_ = dict((k,v) for k,v in enumerate(GENIE_MAP_6))

    cipher = sum((to_bits(genie_ord(c),4) for c in code), ())
    plain  = genie_shuffle(cipher, map_)

    addr  = 0x8000 | from_bits(plain[0:16])
    value = from_bits(plain[16:24])
    print("0x{:04X}\t0x{:02X}".format(addr, value))

def genie_decode_8(code):
    map_ = dict((k,v) for k,v in enumerate(GENIE_MAP_8))

    cipher = sum((to_bits(genie_ord(c),4) for c in code), ())
    plain  = genie_shuffle(cipher, map_)

    addr  = 0x8000 | from_bits(plain[0:16])
    value = from_bits(plain[16:24])
    cmp_  = from_bits(plain[24:32])
    print("0x{:04X}\t0x{:02X}\t0x{:02X}".format(addr, value, cmp_))

def genie_decode(args):
    code = args.code

    if len(code) == 6:
        genie_decode_6(code)
    else:
        genie_decode_8(code)


#---------------------------------------------------------------------
# main
#---------------------------------------------------------------------

def arg_addr15(str_):
    value = int(str_, base=0)
    if not 0x8000 <= value <= 0xFFFF:
        raise argparse.ArgumentTypeError("invalid address: {}".format(str_))
    return value

def arg_u8(str_):
    value = int(str_, base=0)
    if not 0 <= value <= 0xFF:
        raise argparse.ArgumentTypeError("invalid value: {}".format(str_))
    return value

def arg_genie(str_):
    if len(str_) not in (6, 8):
        raise argparse.ArgumentTypeError("invalid game genie code: {}".format(str_))
    if any(c not in GENIE_CHARS for c in str_):
        raise argparse.ArgumentTypeError("invalid game genie code: {}".format(str_))
    return str_

def parse_args():
    ap = argparse.ArgumentParser()
    sps = ap.add_subparsers(dest="cmd")
    sps.required = True

    ap_enc = sps.add_parser("enc")
    ap_enc.set_defaults(func=genie_encode)
    ap_enc.add_argument("addr", type=arg_addr15)
    ap_enc.add_argument("value", type=arg_u8)
    ap_enc.add_argument("cmp_", type=arg_u8, nargs="?")

    ap_dec = sps.add_parser("dec")
    ap_dec.set_defaults(func=genie_decode)
    ap_dec.add_argument("code", type=arg_genie)

    return ap.parse_args()

def main():
    args = parse_args()

    args.func(args)

if __name__ == "__main__": main()