#                       Metasploit NASM Shell Python Mimic
#
# A mimic of metasploit's msf-nasm_shell ruby script, a CLI for converting assembly instructions
# into bytes. Very useful for cases where your code needs to work with assembly but the language isn't
# fully capable or bad at handling assembly on its own. 
#
#                              https://mahemium.blog/


from keystone import *
from rich import print
import sys # For arguments


DEFAULT_ARCHITECTURE = "32" # 32 or 64


def should_use_64_bit():
    result = True if DEFAULT_ARCHITECTURE == "64" else False
    if len(sys.argv) > 1:
        if sys.argv[1] == "64":
            result = True
        elif sys.argv[1] == "32" or  sys.argv[1] == "86":
            result = False
        else:
            print("{}: [red]Invalid architecture '{}'".format(sys.argv[0], sys.argv[1]))
            exit(-1)
    return result


def main():
    is_64bit = should_use_64_bit()
    ks = Ks(KS_ARCH_X86, KS_MODE_32 if not is_64bit else KS_MODE_64)
    
    try:
        while True:
            print("[bold]({}) nasm [/bold]> ".format("x64" if is_64bit else "x86"), end="")
            ops = input().split(";")
            
            offset = 0
            for op in ops:
                op = op.strip()
                try:
                    encoding, count = ks.asm(op.strip())
                except keystone.KsError as e:
                    print("{}: [red]{} [white]-> `{}`".format(sys.argv[0], str(e), op))
                    continue
                    
                b = "".join(["{:02x}".format(x).upper() for x in encoding])
                b = b + " " * ((32 if is_64bit else 16) - len(b))
                f = "{:016x}" if is_64bit else "{:08x}"
                print("[aqua]{}[/aqua]  [white]{}[/white]  [yellow]{}[/yellow]".format(f.format(offset), b, op))
                offset += len(encoding)
                
    except KeyboardInterrupt:
        print("\nExiting...")
        exit(0)
        

if __name__ == "__main__":
    main()