import os
import random
import subprocess
import sys
from pwn import *
from LibcSearcher import *
import inspect

'''
明知道是陷阱，
    为什么还要来。
'''
RED = '\033[91m'
GREEN = '\033[92m'
END = '\033[0m'
YELLOW = '\033[93m'

n2b = lambda x: str(x).encode()
rv = lambda x: p.recv(x)
rl = lambda: p.recvline()
ru = lambda s: p.recvuntil(s)
sd = lambda s: p.send(s)
sl = lambda s: p.sendline(s)
sn = lambda s: sl(n2b(n))
sa = lambda t, s: p.sendafter(t, s)
sla = lambda t, s: p.sendlineafter(t, s)
sna = lambda t, n: sla(t, n2b(n))
ia = lambda: p.interactive()
rop = lambda r: flat([p64(x) for x in r])
uu64 = lambda data: u64(data.ljust(8, b'\x00'))
rall = lambda : p.recvall()

## Initialize world
def libset(libc_val):
    """
    Set libc for getting offset
    """
    global libc
    libc = ELF(libc_val)

def setup(p_val):
    """
    Set program and start
    """
    global p
    global elf
    p = process(p_val)
    elf = ELF(p_val)

def set(p_val):
    """
    Set program without starting
    """
    global elf
    elf = ELF(p_val)

def rsetup(mip, mport):
    """
    Set up remote connection (remote setup)
    """
    if args.P:
        global p
        p = remote(mip, mport)

## Receive world
def tet():
    """
    Test receiving data for a line 
    """
    p = globals()['p']
    r = ru('\n')
    print('\n----------------\n', 'add', 'is >>> ', r, '\n---------------')
    return r

def getx64(i, j):
    """
    Get 64bit and unpack it with debug info
    """
    if i != 0:
        r = (ru('\n'))[i:j]
        dp('getx64', r)
        r = u64(r.ljust(8, b'\0'))
        print('\n----------------\n', 'add', 'is >>> ', hex(r), '\n---------------')
        return r
    else:
        r = (ru('\n'))[:j]
        dp('getx64', r)
        r = u64(r.ljust(8, b'\0'))
        print('\n----------------\n', 'add', 'is >>> ', hex(r), '\n---------------')
        return r

def getx32(i, j):
    """
    Get 32bit and unpack it with debug info
    """
    if i != 0:
        r = (ru('\n'))[i:j]
        dp('getx32', r)
        r = u32(r.ljust(4, b'\0'))
        print('\n----------------\n', 'add', 'is >>> ', hex(r), '\n---------------')
        return r
    else:
        r = (ru('\n'))[:j]
        dp('getx32', r)
        r = u32(r.ljust(4, b'\0'))
        print('\n----------------\n', 'add', 'is >>> ', hex(r), '\n---------------')
        return r

def getx(i, j):
    """
    Get hex and unpack it with debug info
    """
    if i != 0:
        r = (ru('\n'))[i:j]
        dp('getx', r)
        r = int(r, 16)
        print('\n----------------\n', 'add', 'is >>> ', hex(r), '\n---------------')
        return r
    else:
        r = (ru('\n'))[:j]
        dp('getx', r)
        r = int(r, 16)
        print('\n----------------\n', 'add', 'is >>> ', hex(r), '\n---------------')
        return r

def getd(i, j):
    """
    Get decimal and unpack it with debug info
    """
    if i != 0:
        r = (ru('\n'))[i:j]
        dp('getd', r)
        r = int(r, 10)
        print('\n----------------\n', 'add', 'is >>> ', hex(r), '\n---------------')
        return r
    else:
        r = (ru('\n'))[:j]
        dp('getd', r)
        r = int(r, 10)
        print('\n----------------\n', 'add', 'is >>> ', hex(r), '\n---------------')
        return r

def tryre():
    """
    Try re for check if the connection is still alive
    """
    return p.recvrepeat(2)

def close():
    """
    Close connection
    """
    p.close()

'''
只攻不防，
    天下无双—————
        魔刀千刃。
'''
## Calculate world

def getbase(add, defname, *args):
    """
    Get base from the real address you got
    """
    base = add - libc.symbols[defname]
    for num in args:
        base -= num
    print('\nloading...')
    print('\n----------------\nget!your base is >>> ', hex(base), '\n--------------')
    return base

def evgdb(*argv):
    """
    Set gdb (evil-gdb)
    """
    p = globals()['p']
    context.terminal = ['alacritty', '-e']
    # Modify terminal as per your environment
    # Replace 'alacritty' with the result of running 'echo $TERM'
    if args.G:
        if(len(argv)==0):
            gdb.attach(p)
        else:
            gdb.attach(p, argv[0])

def symoff(defname, *args):
    """
    Calculate or set symblol's offset
    """
    if(len(args)>0):
        ba = args[0]
        print('\n----------------\nyour ', defname, 'offset is >>> ', hex(libc.sym[defname]), '\n---------------')
        print('\n----------------\nyour ', defname, 'is in >>> ', hex(ba+libc.sym[defname]), '\n---------------')
        return libc.sym[defname]+ba
    else:
        print('\n---------------\nyour ', defname, 'offset is >>> ', hex(libc.sym[defname]), '\n---------------')
        return libc.sym[defname]

def gotadd(defname, *args):
    """
    Get got's address
    """
    dpx(f"{defname} got's add", elf.got[defname])
    if (len(args) > 0):
        return elf.got[defname]+args[0]  # Handle PIE
    return elf.got[defname]

def pltadd(defname, *args):
    """
    Get plt's address
    """
    dpx(f"{defname} plt's add", elf.got[defname])
    if (len(args) > 0):
        return elf.plt[defname]+args[0]  # Handle PIE
    return elf.plt[defname]

def symadd(defname, *args):
    """
    Get sym's address
    """
    dpx(f"{defname} sym's add", elf.got[defname])
    if (len(args) > 0):
        return elf.sym[defname]+args[0]  # Handle PIE
    return elf.sym[defname]

def dp(name, data):
    """
    Print data with name
    """
    print('\n---------------\nyour ', name, ' is >>> ', (data), '\n---------------')

def dpx(name, data):
    """
    Print hex data with name
    """
    print('\n---------------\nyour ', name, ' is >>> ', hex(data), '\n---------------')

def d(data):
    """
    Print data without name you input and get name from your local variable
    """
    frame = inspect.currentframe().f_back
    locals_dict = frame.f_locals

    for name, value in locals_dict.items():
        if value is data:
            print('\n---------------\nyour', name, 'is >>>', data, '\n---------------')
            break

def dx(data):
    """
    Hex print data without name you input and get name from your local variable
    """
    frame = inspect.currentframe().f_back
    locals_dict = frame.f_locals

    for name, value in locals_dict.items():
        if value is data:
            print('\n---------------\nyour', name, 'is >>>', hex(data), '\n---------------')
            break

'''
因为，   
    我有想要保护的人。
'''

## Library exploration world

def rlibset(defname, add):
    """
    Set remote libc
    """
    global rlibc
    rlibc = LibcSearcher(defname, add)

def rgetbase(add, defname, *args):
    """
    Get remote libc base from the real address you got
    """
    base = add - rlibc.dump(defname)
    for num in args:
        base -= num
    print('\nloading...')
    print('\n----------------\nget!your base is >>> ', hex(base), '\n--------------')
    return base

def rsymoff(defname, *args):
    """
    Calculate or set remote symblol's offset
    """
    if(len(args)>0):
        ba = args[0]
        print('\n----------------\nyour ', defname, 'offset is >>> ', hex(rlibc.dump(defname)), '\n---------------')
        print('\n----------------\nyour ', defname, 'is in >>> ', hex(ba+rlibc.dump(defname)), '\n---------------')
        return rlibc.dump(defname)+ba
    else:
        print('\n---------------\nyour ', defname, 'offset is >>> ', hex(rlibc.dump(defname)), '\n---------------')
        return rlibc.dump(defname)

# Attack world

def fmt(offset, begin, end, size, written):
    """
    Format string exploit
    """
    payload = fmtstr_payload(offset, {begin: end}, write_size=size, numbwritten=written)
    return payload

'''
    offset（int） - 您控制的第一个格式化程序的偏移量
    字典（dict） - 被写入地址对应->写入的数据，可多个对应{addr: value, addr2: value2}
    numbwritten（int） - printf函数已写入的字节数
    write_size（str） - 必须是byte，short或int。告诉您是否要逐字节写入，短按short或int（hhn，hn或n）
    注意还有ln，lln等用法
'''
def re(time):
    """
    Try to receive
    """
    p.recvS(timeout=time)  # Used for brute-forcing
def clo():
    """
    Close connection
    """
    p.close()

'''
while True:
    ...
    if re:
        print('pwned!get your flag here:', re)
        exit(0)
    p.close()以udi
'''
got = gotadd
plt = pltadd
off = symoff

def execute_script(script_path, mouse=None):
    try:
        # 在脚本文件中插入赋值语句
        if mouse is not None:
            with open(script_path, 'r+') as f:
                content = f.read()
                f.seek(0, 0)
                f.write(f"mouse = {mouse}\n" + content)
        
        process = subprocess.Popen(['python3', script_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE, text=True)
        sleep(0.4)
        output, errors = process.communicate(input='l')  
        
        return_code = process.returncode
        
        return output, errors, return_code
    finally:
        if mouse is not None:
            with open(script_path, 'r+') as f:
                lines = f.readlines()
                f.seek(0)
                for line in lines:
                    if not line.startswith(f"mouse = {mouse}"):
                        f.write(line)
                f.truncate()


import time

def cyccat(your_exp, offset_you_know, overflow_key):
    """
    get offset of stackoverflow autoly

    usage:
    for first, you must code a "mouse" as a variable in the input in your exploit script for testing the offset
    like:
    sd(mouse)

    your_exp : your exploit script
    offset_you_know : the offset can make overflow you know
    overflow_key : the key words you know can make overflow, if it's 0, it means EOF
    """
    min_range = 0
    max_range = offset_you_know
    start_time = time.time()
    rotating_symbols = ['.', 'o', '0', 'O']
    rotating_index = 0
    last_print_time = start_time
    while min_range <= max_range:
        guess = random.randint(min_range, max_range)
        mouse = cyclic(guess)
        out, err, return_code = execute_script(your_exp, mouse=mouse)

        if overflow_key == 0:
            condition_met = "Got EOF" in out
        else:
            condition_met = overflow_key in out

        if condition_met:
            max_range = guess - 1
        else:
            min_range = guess + 1
        
        # 检查是否到达打印时间间隔
        current_time = time.time()
        if current_time - last_print_time >= 0.1:  # 设置时间间隔为0.1秒
            # 打印当前进度
            rotating_symbol = rotating_symbols[rotating_index % len(rotating_symbols)]
            rotating_index += 1
            progress = (max_range - min_range) / offset_you_know * 100
            print(f"{YELLOW}[{rotating_symbol}]{END} Progress: {100 - 0.05 - progress:.2f}%\r", end='', flush=True)
            last_print_time = current_time

    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"{YELLOW}[{rotating_symbol}]{END} Progress: 100.00%\r", end='', flush=True)
    print(f"\n[!] Time elapsed: {elapsed_time:.2f} seconds")
    if min_range-1 == offset_you_know:
        print(f"{RED}[-]{END} Not found! Please check your overflow key or offset, and try again")
    else:
        print(f"{GREEN}[+]{END} Found! The appropriate payload size: {min_range-1} (just covering the return address)")


def congra():
    """
    a congratulation message
    """
    print(f"{GREEN}[+]{END} Pwned! Go ahead")
