compilefunc

#!/usr/bin/env python

compilefunc_usage = (
'''usage: compilefunc [options] SOURCE-FILE FUNCTION [FUNCTION...]
Generate assembler output on stdout for specified functions.
Supports the CC, CFLAGS, CXX, CXXFLAGS, LDFLAGS and GDB environment variables.
Unknown options are passed to the compiler.
Invokes gcc or g++ depending on the file extension.

    $ CFLAGS='-O0' compilefunc --do-link var-arg.c main
       0x000000000040067b <+0>: push   %rbp
       0x000000000040067c <+1>: mov    %rsp,%rbp
       0x000000000040067f <+4>: movsd  0xc9(%rip),%xmm0        # 0x400750
       0x0000000000400687 <+12>:    mov    $0x1,%esi
       0x000000000040068c <+17>:    mov    $0x40074a,%edi
       0x0000000000400691 <+22>:    mov    $0x1,%eax
       0x0000000000400696 <+27>:    callq  0x400536 <f>
       0x000000000040069b <+32>:    mov    $0x0,%eax
       0x00000000004006a0 <+37>:    pop    %rbp
       0x00000000004006a1 <+38>:    retq   
    $ CFLAGS='-Os' compilefunc --do-link var-arg.c main
       0x0000000000400440 <+0>: push   %rax
       0x0000000000400441 <+1>: mov    $0x1,%esi
       0x0000000000400446 <+6>: mov    $0x400674,%edi
       0x000000000040044b <+11>:    mov    $0x1,%al
       0x000000000040044d <+13>:    movsd  0x22b(%rip),%xmm0        # 0x400680
       0x0000000000400455 <+21>:    callq  0x400556 <f>
       0x000000000040045a <+26>:    xor    %eax,%eax
       0x000000000040045c <+28>:    pop    %rdx
       0x000000000040045d <+29>:    retq   
    $ CFLAGS='-O3' compilefunc --do-link var-arg.c main
       0x0000000000400440 <+0>: sub    $0x8,%rsp
       0x0000000000400444 <+4>: mov    $0x1,%esi
       0x0000000000400449 <+9>: mov    $0x400694,%edi
       0x000000000040044e <+14>:    movsd  0x24a(%rip),%xmm0        # 0x4006a0
       0x0000000000400456 <+22>:    mov    $0x1,%eax
       0x000000000040045b <+27>:    callq  0x400570 <f>
       0x0000000000400460 <+32>:    xor    %eax,%eax
       0x0000000000400462 <+34>:    add    $0x8,%rsp
       0x0000000000400466 <+38>:    retq  

NOTE: unless --do-link (or -l) is passed, global variable and function references
will appear to reference the wrong address.
''')

gdbdis_usage = (
'''usage: gdbdis EXE FUNC
Disassemble a funtion from an executable with GDB
        $ ./gdbdis /bin/true "'realloc@plt'" main,+6
        Dump of assembler code for function realloc@plt:
           0x0000000000401240 <+0>: jmpq   *0x204ee2(%rip)        # 0x606128 <realloc@got.plt>
           0x0000000000401246 <+6>: pushq  $0x22
           0x000000000040124b <+11>:    jmpq   0x401010
        End of assembler dump.
        Dump of assembler code from 0x401320 to 0x401326:
           0x0000000000401320 <main+0>: cmp    $0x2,%edi
           0x0000000000401323 <main+3>: push   %rbx
           0x0000000000401324 <main+4>: je     0x40132d <main+13>
        End of assembler dump.

With a single argument, the function surrounding that address is dumped.
Two arguments (separated by a comma) are taken as a range of memory to dump,
  in the form of "start,end", or "start,+length".

Note that the address is interpreted as an expression, not as a location
like in the "break" command.
So, for example, if you want to disassemble function bar in file foo.c
you must type "disassemble 'foo.c'::bar" and not "disassemble foo.c:bar".

Environment Variables: GDB
''')

gdbx_usage = (
'''usage: gdbx EXE ADDR FORMAT
Examine memory with GDB from an executable
        $ gdbx /bin/true 0x4039e4 s
        0x4039e4 <__dso_handle+732>:   "GNU coreutils"

        $ GDB=arm-none-eabi-gdb gdbx cortex-m.elf 0 2a
        0x0 <__isr_vector>: 0x20001000  0x185 <Reset_Handler> 

ADDR is an expression for the memory address to examine.
FORMAT is a repeat count followed by a format letter and a size letter.
Format letters are o(octal), x(hex), d(decimal), u(unsigned decimal),
  t(binary), f(float), a(address), i(instruction), c(char), s(string)
  and z(hex, zero padded on the left).
Size letters are b(byte), h(halfword), w(word), g(giant, 8 bytes).

Environment Variables: GDB
''')

import sys
import os
import subprocess
import errno
import argparse
import tempfile
import logging
from dataclasses import dataclass

main_function_map = {}

def main_function(func):
    global main_function_map
    main_function_map[func.__name__.replace('_','-')] = func
    return func

def main_function_dispatch(name, args):
    try:
        f = main_function_map[name]
    except KeyError:
        sys.stderr.write('%s is not a valid command name\n' % (name,))
        sys.exit(2)
    f(args)

def program_name():
    return os.path.basename(sys.argv[0])

def gdb_batch(gdb_cmd, command_list):
    # -> output_str
    gdb_input = tempfile.NamedTemporaryFile(mode='w+')
    gdb_input.write('\n'.join(command_list))
    gdb_input.seek(0)

    (out_file, err_file) = (tempfile.TemporaryFile(mode='w+'),
                            tempfile.TemporaryFile(mode='w+'))
    r = subprocess.call([gdb_cmd, '--batch', '-x', gdb_input.name],
                        stdout=out_file, stderr=err_file)
    err_file.seek(0)
    if r != 0:
        raise ValueError(err_file.read())

    # NOTE: gdb exits with 0 even on errors ...
    # Let's hope these GDB error message strings are stable ..
    
    # Ignore the following:
    # 'warning: the debug information found in "/usr/lib/debug//usr/bin/true.debug" does not match "/usr/bin/true" (CRC mismatch).\n',
    err_lines = [ x for x in err_file.readlines() if x.strip() and not x.startswith('warning: ') ]
    if err_lines:
        # err_lines[0] would be 't.gdb:2: Error in sourced command file:', skip it
        err_lines = err_lines[1:]
        out = []
        out.append('%s input:\n' % (gdb_cmd,))
        gdb_input.seek(0)
        out.extend(['  ' + x for x in gdb_input.readlines()])
        out.append('\n%s stderr:\n' % (gdb_cmd,))
        out.extend(['  ' + x for x in err_lines])
        raise ValueError(''.join(out))
    out_file.seek(0)
    return out_file.read()

def gdb_prepare_disassemble_cmd(gdb_cmd, exe, options):
    # -> (batch_cmds, dis_options)
    dis_options = ''
    if options.raw_instructions:
        dis_options += 'r'
    if options.source:
        dis_options += 's'
    elif options.source_centric:
        dis_options += 'm'
    # prefix with '/'
    if dis_options:
        dis_options = '/' + dis_options

    batch_cmds = [
        ('file %s' % (exe,)),
    ]
    if getattr(options, 'intel_syntax', None):
        batch_cmds.append('set disassembly-flavor intel')
    if getattr(options, 'demangle', None):
        batch_cmds.append('set print asm-demangle')
    if getattr(options, 'starti', None):
        batch_cmds.append('starti')
    return batch_cmds, dis_options

def gdb_disassemble_functions(gdb_cmd, exe, functions, options):
    # -> output_str
    cmd, dis_options = gdb_prepare_disassemble_cmd(gdb_cmd, exe, options)
    for i in functions:
        if not options.no_special_symbol_quoting:
            #"printf@plt" -> "'printf@plt'"
            if i.endswith('@plt'):
                i = "'%s'" % (i,)
        cmd.append('disassemble %s %s' % (dis_options, i))
    o = gdb_batch(gdb_cmd, cmd)

    if not options.no_header:
        return o

    (start, end) = (-1, 0)
    out = []
    while 1:
        start = o.find('Dump of assembler code for function ', start + 1)
        if start == -1:
            break
        start = o.find('\n', start) + 1
        end = o.find('End of assembler dump.\n', start)
        assert(end != -1)
        out.append(o[start:end])
    return ''.join(out)

def gdb_disassemble_range(gdb_cmd, exe, range_start, range_end, options):
    # -> output_str
    cmd, dis_options = gdb_prepare_disassemble_cmd(gdb_cmd, exe, options)
    cmd.append('disassemble %s %s,%s' % (dis_options, range_start, range_end))
    o = gdb_batch(gdb_cmd, cmd)
    return o

def gdb_command_name():
    cc = os.environ.get('CC', 'gcc')
    gdb = os.environ.get('GDB')
    if gdb is None:
        # 'arm-linux-androideabi-gcc' -> 'arm-linux-androideabi-gdb'
        try:
            gdb = cc[:cc.rindex('-')+1]+'gdb'
        except ValueError:
            gdb = 'gdb'
    return gdb

def readelf_command_name():
    cc = os.environ.get('CC', 'gcc')
    readelf = os.environ.get('READELF')
    if readelf is None:
        # 'arm-linux-androideabi-gcc' -> 'arm-linux-androideabi-readelf'
        try:
            readelf = cc[:cc.rindex('-')+1]+'readelf'
        except ValueError:
            readelf = 'readelf'
    return readelf

@main_function
def compilefunc(args):
    ap = argparse.ArgumentParser(description='generate asm for the specified functions')
    ap.add_argument('source_file', metavar='SOURCE-FILE', help='C or C++ source file')
    ap.add_argument('functions', metavar='FUNCTION', nargs='+', help='function to generate asm for')
    ap.add_argument('-l', '--do-link', action='store_true', help='link and not just compile the code')
    parser_add_disassemble_arguments(ap)
    (args, unknown_args) = ap.parse_known_args(args)
    if len(args.functions) < 2:
        args.no_header = True

    t = args.source_file
    if t.endswith('.cpp') or t.endswith('.cc'):
        (env_cc, default_cc, env_cflags) = ('CXX', 'g++', 'CXXFLAGS')
    else:
        (env_cc, default_cc, env_cflags) = ('CC', 'gcc', 'CFLAGS')
    cc = os.environ.get(env_cc, default_cc)
    cflags = os.environ.get(env_cflags, '')
    cflags = [ x for x in cflags.split() if x ]
    gdb = gdb_command_name()

    (fout_fd, fout) = tempfile.mkstemp()
    t = ['-g']
    if args.do_link:
        ldflags = os.environ.get('LDFLAGS', '')
        ldflags = [ x for x in ldflags.split() if x ]
        t.extend([ args.source_file ] + ldflags)
    else:
        t.extend([ '-c', args.source_file ])
    o = None
    try:
        cmd = [cc] + cflags  + unknown_args + t + ['-o', fout]
        logging.debug('compile: %s' % (cmd,))
        r = subprocess.call(cmd)
        if r == 0:
            try:
                o = gdb_disassemble_functions(gdb, fout, args.functions, args)
            except ValueError:
                sys.exit(3)
    finally:
        os.unlink(fout)
    if o is not None:
        sys.stdout.write(o)

def parser_add_disassemble_arguments(ap):
    ap.add_argument('-r', '--raw-instructions', action='store_true', help='show raw instruction opcodes')
    ap.add_argument('-s', '--source', action='store_true', help='include source lines'),
    ap.add_argument('-m', '--source-centric', action='store_true', help='include `disassemble /m` source lines'),
    ap.add_argument('--no-header', action='store_true', 
                    help="don't show \"Dump of ...\" headers and footers even when disassembling multiple functions")
    # --no-special-symbol-quoting: created so that by default,
    # you don't need to do e.g. "'printf@plt'" from the shell to pass in PLT functions
    ap.add_argument('--no-special-symbol-quoting', action='store_true', default=False,
                    help='disable auto quoiting of symbols that ends with @plt')
    ap.add_argument('--demangle', action='store_true', default=False,
                    help='demangle C++ names in disassembly listings')
    ap.add_argument('--intel-syntax', action='store_true', default=False,
                    help='disassemble x86 instrucitons using Intel instead of AT&T syntax')
    ap.add_argument('--starti', action='store_true', help='Start the debugged program stopping at the first instruction')
    return ap

@main_function
def gdbdis(args):
    ap = argparse.ArgumentParser(description='disassemble functions from EXE')
    ap.add_argument('executable', metavar='EXECUTABLE', help='executable file to disassemble')
    ap.add_argument('functions', metavar='FUNCTION', nargs='+', help='function name or address to disassemble')
    parser_add_disassemble_arguments(ap)
    (args, unknown_args) = ap.parse_known_args(args)

    if len(args.functions) == 1:
        args.no_header = True

    gdb = gdb_command_name()
    try:
        o = gdb_disassemble_functions(gdb, args.executable, args.functions, args)
    except ValueError as e:
        sys.stderr.write(str(e))
        sys.exit(3)
    sys.stdout.write(o)

# functions that parse `readelf --program-headers --wide` output

@dataclass
class ElfSegment:
    '''`readelf --program-header` data for an ELF segment'''
    type: str
    offset: int
    virt_addr: int
    phys_addr: int
    file_size: int
    mem_size: int
    flag: set
    align: int

def program_header_load_line_parse(line):
    f = line.split()
    # Handle 'RW', 'R E' etc in the Flg field
    flag_str = ''.join(x for x in f[6:-1]) # start after MemSiz exclude Align
    return ElfSegment(type=f[0],
                   offset=int(f[1], base=16),
                   virt_addr=int(f[2], base=16),
                   phys_addr=int(f[3], base=16), 
                   file_size=int(f[4], base=16), 
                   mem_size=int(f[5], base=16), 
                   flag=set(flag_str),
                   align=int(f[-1], base=16))

def program_headers_parse_executable_segment_file_size(lines):
    sample = '''Program Headers:
      Type           Offset   VirtAddr           PhysAddr           FileSiz  MemSiz   Flg Align
      PHDR           0x000040 0x0000000000000040 0x0000000000000040 0x0002d8 0x0002d8 R   0x8
      INTERP         0x000318 0x0000000000000318 0x0000000000000318 0x00001c 0x00001c R   0x1
          [Requesting program interpreter: /lib64/ld-linux-x86-64.so.2]
      LOAD           0x000000 0x0000000000000000 0x0000000000000000 0x001040 0x001040 R   0x1000
      LOAD           0x002000 0x0000000000002000 0x0000000000002000 0x002c51 0x002c51 R E 0x1000'''
    #lines = sample.split('\n')
    segment = None
    for line in lines:
        #print('line: %r' % (line,))
        if not line.strip().startswith('LOAD '):
            continue
        segment = program_header_load_line_parse(line)
        if 'E' in segment.flag:
            break
    if segment is None:
        raise ValueError('executable segment not found')
    #assert(segment.file_size == 0x002c51)
    #assert(segment.flag == set('re'))
    return segment.file_size 

def executable_segment_size(exe_path):
    readelf = subprocess.run([readelf_command_name(), '--program-headers', '--wide', exe_path], capture_output=True, check=True,
                            text=True, encoding='ascii')
    size = program_headers_parse_executable_segment_file_size(readelf.stdout.split('\n'))
    return size

@dataclass
class ProcMapping:
    '''data from GDB's `info proc mappigns` output'''
    start_addr: int
    end_addr: int
    size: int
    offset: int
    perms: set
    objfile: str

def info_proc_mappings_executable_segment_start(lines):
    sample = '''(gdb) info proc mappings 
    process 2182810
    Mapped address spaces:

              Start Addr           End Addr       Size     Offset  Perms  objfile
          0x555555554000     0x555555555000     0x1000        0x0  r--p   /home/scottt/work/pwn.college/babyrev.2022fall/babyrev_level16.0
          0x555555555000     0x555555557000     0x2000     0x1000  r-xp   /home/scottt/work/pwn.college/babyrev.2022fall/babyrev_level16.0
          0x555555557000     0x555555558000     0x1000     0x3000  r--p   /home/scottt/work/pwn.college/babyrev.2022fall/babyrev_level16.0
          0x555555558000     0x55555555a000     0x2000     0x3000  rw-p   /home/scottt/work/pwn.college/babyrev.2022fall/babyrev_level16.0
          0x7ffff7fc0000     0x7ffff7fc4000     0x4000        0x0  r--p   [vvar]
          0x7ffff7fc4000     0x7ffff7fc6000     0x2000        0x0  r-xp   [vdso]
          0x7ffff7fc6000     0x7ffff7fc8000     0x2000        0x0  r--p   /usr/lib64/ld-linux-x86-64.so.2
          0x7ffff7fc8000     0x7ffff7fef000    0x27000     0x2000  r-xp   /usr/lib64/ld-linux-x86-64.so.2
          0x7ffff7fef000     0x7ffff7ffa000     0xb000    0x29000  r--p   /usr/lib64/ld-linux-x86-64.so.2
          0x7ffff7ffb000     0x7ffff7fff000     0x4000    0x34000  rw-p   /usr/lib64/ld-linux-x86-64.so.2
          0x7ffffffdc000     0x7ffffffff000    0x23000        0x0  rw-p   [stack]
      0xffffffffff600000 0xffffffffff601000     0x1000        0x0  --xp   [vsyscall]'''
    #lines = sample.split('\n')
    for line in lines:
        f = line.split()
        #print('info_proc_mappings: f: %r' % (f,))
        if len(f) != 6:
          continue
        try:
          start_addr = int(f[0], base=16)
          end_addr = int(f[1], base=16)
        except ValueError:           
            continue
        #print('start_addr: 0x%x' % (start_addr,))
        # perms: handle `rw-p` where '-' is for missing permissions and 'p' I assume just marks the end
        perms = set(f[4])
        try:
            perms.remove('-')
            perms.remove('p')
        except KeyError:
            pass
        m = ProcMapping(start_addr=start_addr,
                           end_addr=end_addr,
                           size=int(f[2], base=16),
                           offset=int(f[3], base=16),
                           perms=perms,
                           objfile=f[5])
        if 'x' in m.perms:
            #assert(m.start_addr == 0x555555555000)
            #assert(m.perms == set('rx'))
            return m.start_addr
    else:
        raise ValueError('executable segment not found')

def gdb_executable_segment_start(gdb_cmd, options):
    # Requirement: work on stripped PIEs, thus using `starti` instead of `start` etc
    cmd = [ 'file %s' % (options.executable,),
           'starti',
           'info proc mappings'
    ]
    o = gdb_batch(gdb_cmd, cmd)
    return info_proc_mappings_executable_segment_start(o.split('\n'))

@main_function
def gdbdis_exe(args):
    # BUG: garbage output from `starti`: `Program stopped.` and  `... in _start() from /lib64/...`
    #
    #$ gdbdis-exe --intel-syntax /bin/true 2>/dev/null | head
    #Program stopped.
    #0x00007ffff7fe3d10 in _start () from /lib64/ld-linux-x86-64.so.2
    #Dump of assembler code from 0x555555556000 to 0x555555558c51:
    #   0x0000555555556000 <_init+0>:	endbr64
    #   0x0000555555556004 <_init+4>:	sub    rsp,0x8
    #   0x0000555555556008 <_init+8>:	mov    rax,QWORD PTR [rip+0x5fc1]        # 0x55555555bfd0
    #   0x000055555555600f <_init+15>:	test   rax,rax
    #   0x0000555555556012 <_init+18>:	je     0x555555556016 <_init+22>
    #   0x0000555555556014 <_init+20>:	call   rax

    ap = argparse.ArgumentParser(description='disassemble executable segment from EXE')
    ap.add_argument('executable', metavar='EXECUTABLE', help='executable file to disassemble')
    parser_add_disassemble_arguments(ap)
    (args, unknown_args) = ap.parse_known_args(args)
    args.starti = True
    gdb = gdb_command_name()
    # TODO: don't load exe into gdb twice just to find the executable segment load address for PIEs
    range_start = gdb_executable_segment_start(gdb, args)
    range_end = '+0x%x' % (executable_segment_size(args.executable),)

    try:
        o = gdb_disassemble_range(gdb, args.executable, range_start, range_end, args)
    except ValueError as e:
        msg = getattr(e, 'message', None)
        if msg is None:
            msg = str(e)
        sys.stderr.write(msg)
        sys.exit(3)
    try:
        sys.stdout.write(o)
    except BrokenPipeError:
        pass

@main_function
def gdbx(args):
    ap = argparse.ArgumentParser(description='examine memory from EXE')
    ap.add_argument('executable', metavar='EXECUTABLE', help='executable file to examine')
    ap.add_argument('address', metavar='ADDRESS', help='symbol name or address to examine')
    ap.add_argument('format', metavar='FORMAT', help='format for GDB like 4i, 8hx')
    ap.add_argument('--no-special-symbol-quoting', action='store_true', default=False,
                    help='disable auto quoiting of symbols that ends with @plt')
    (args, unknown_args) = ap.parse_known_args(args)


    if not args.no_special_symbol_quoting:
        if args.address.endswith('@plt'):
            args.address = "'%s'" % (args.address,)

    gdb = gdb_command_name()
    cmd = [
        ('file %s' % (args.executable,)),
        ('x/%(format)s %(address)s' % dict(format=args.format, address=args.address)),
    ]
    try:
        o = gdb_batch(gdb, cmd)
    except ValueError as e:
        sys.stderr.write(e.message)
        sys.exit(3)
    sys.stdout.write(o)

@main_function
def gdb_info_scope(args):
    ap = argparse.ArgumentParser(description='show debug info for function')
    ap.add_argument('executable', metavar='EXECUTABLE', help='executable file to examine')
    ap.add_argument('functions', metavar='FUNCTION', nargs='+', help='function name or address to disassemble')
    ap.add_argument('--no-special-symbol-quoting', action='store_true', default=False,
                    help='disable auto quoiting of symbols that ends with @plt')
    (args, unknown_args) = ap.parse_known_args(args)

    gdb = gdb_command_name()
    cmd = [
        'file '+ args.executable,
    ]

    for i in args.functions:
        if not args.no_special_symbol_quoting:
            if i.endswith('@plt'):
                i = "'%s'" % (i,)
        cmd.append('info scope ' + i)
    try:
        o = gdb_batch(gdb, cmd)
    except ValueError as e:
        sys.stderr.write(e.message)
        sys.exit(3)
    sys.stdout.write(o)

def gdb_info_cmd(description, info_cmd, args):
    ap = argparse.ArgumentParser(description=description)
    ap.add_argument('executable', metavar='EXECUTABLE', help='executable file to examine')
    (args, unknown_args) = ap.parse_known_args(args)

    gdb = gdb_command_name()
    cmd = [
        'file '+ args.executable,
    ]
    if unknown_args:
        cmd.append('info %s %s' % (info_cmd, unknown_args[0]))
    else:
        cmd.append('info %s' % (info_cmd, ))

    try:
        o = gdb_batch(gdb, cmd)
    except ValueError as e:
        sys.stderr.write(e.message)
        sys.exit(3)
    sys.stdout.write(o)

@main_function
def gdb_info_functions(args):
    gdb_info_cmd('list functions from EXE', 'functions', args)

@main_function
def gdb_info_variables(args):
    gdb_info_cmd('list variables from EXE', 'variables', args)

@main_function
def gdb_info_types(args):
    gdb_info_cmd('list types from EXE', 'types', args)

@main_function
def gdb_info_line(args):
    gdb_info_cmd('source and memory addresses for a source location from EXE', 'line', args)

if __name__ == '__main__':
    main_function_dispatch(program_name(), sys.argv[1:])