'''
The challenge is straightforward. The server sends a base64 encoded ELF file over
the network. For all binaries, there is a buffer at offset 0xaa8 
in the binary which contains some data (printable ASCII). The flag decryption
function is at 0x7da which loads some values onto the stack and uses them as 
offsets into this data array, stitching them together to validate given input. 
We simply emulate program behavior to find the flag. Pseudocode:

data = ...
index_arr = ..
input = ...
for i in range(len_of_flag): # len_of_flag is operand at 0x954
    if input[i] != data[index_arr[i]]:
        print("Failed!")
        exit(0)
print("Good job!")

This pattern remains the same across all binaries -> including the function 
offset,  the data buffer address, etc. A good approach would be to solve the first 
few problems by hardcoding these values in the automation script and generalising
as and when the server sends a file which is different.
'''

import binaryninja

from pwn import remote
import base64

conn = remote('pwn.byteband.it', 6000)

file_counter = 0
while True:
    # Receive file from network as a base64 encoded string
    elf_base64 = conn.recvline()
    chal = elf_base64.strip()
    elf_base64_decoded = base64.b64decode(chal)

    # Write file onto disk
    chal_name = 'chal_' + str(file_counter)
    with open(chal_name, 'wb') as f:
        f.write(elf_base64_decoded)

    # Load file in BN and perform initial analysis
    bv = binaryninja.BinaryViewType.get_view_of_file(chal_name)
    bv.update_analysis_and_wait()

    # There is a buffer at 0xaa8 which has 30 bytes, read that into data
    data = bv.read(0xaa8, 30).decode('utf-8')

    # Function which does flag computation is at 0x7da
    solver_func = bv.get_function_at(0x7da)

    # Get length of flag from this instruction (remains same for all binaries)
    len_inst = solver_func.get_low_level_il_at(0x954)
    flag_length = len_inst.operands[1].constant

    flag = ""

    # At this address, offsets are pushed into variables on the stack which are
    # used to index into the array.
    # Eg. (sequence of) mov dword [var_xx], 0x12
    # TODO: Use instruction sequence instead of this
    base_addr = 0x7f4
    counter = 0
    while counter < flag_length:
        if counter <= 24:
            inst_addr = base_addr + counter * 10
            inst = solver_func.get_low_level_il_at(inst_addr)
        else:
            inst_addr = base_addr + 240 + (counter - 24) * 7
            inst = solver_func.get_low_level_il_at(inst_addr)

        counter += 1

        # Get the offset and index into data to add character to flag
        offset = inst.operands[1].constant
        flag += data[offset]

    # Send flag
    conn.sendline(flag)
    print("[+] Solved for " + str(file_counter) + " : " + flag)
    file_counter += 1

'''
The idea is simple.

1. TEST_FLAG = ""
2. Hook instruction with compare (flag_byte, input_byte)
2. Run Emulation
3. Check failure_byte (flag_byte at which comparison fails)
4. TEST_FLAG += failure_byte
5. Goto 2
6. We now haz flag!
'''
import sys
from qiling import *

# Suppress qiling syscall emulation debug output
f = open('/dev/null', 'w')
sys.stderr = f

base_path = "/home/chinmay_dd/ByteBanditsCTF/"

# Solving only one challenge for now
# conn = remote('pwn.byteband.it', 6000)
file_counter = 0

# base_flag will eventually contain the flag
base_flag = []
# Flag length collected from EAX value at 96b
length = 0
# Count of how many times the compare instruction ran
count = 0

# Compare function hook
# At the compare instruction, check value in dl which is what is being 
# compared with our input. Add it to global_flag
def h_c(ql):
    global count
    global base_flag
    count += 1
    base_flag += chr(ql.uc.reg_read(UC_X86_REG_AL))

# Instruction hook which will tell us length of flag
def h_d(ql):
    global length
    length = ql.uc.reg_read(UC_X86_REG_EAX)

# Simple class provided with Qiling which will model a stdin/stdout pipe
class MyPipe():
    def __init__(self):
        self.buf = b''
        self.offset = 0

    def write(self, s):
        self.buf += s

    def read(self, size):
        if size <= len(self.buf):
            ret = self.buf[self.offset:self.offset + size]
            self.buf = self.buf[self.offset + size:]
        else:
            ret = self.buf
            self.buf = ''
        return ret

    def fileno(self):
        return 0

    def show(self):
        pass

    def clear(self):
        pass

    def flush(self):
        pass

    # Custom added. Might be wrong, but works.
    def lseek(self, pos, how):
        pass
        if how == 0:
            self.offset = pos
            return pos
        elif how == 1:
            return self.offset
        elif how == 2:
            pass

    def close(self):
        self.outpipe.close()

    def fstat(self):
        return os.fstat(sys.stdin.fileno())

# Function to solve one problem
def solve_one(elf_path):
    global base_flag
    next_flag = ""

    def run_one_round(test_flag):
        # Create two pipes
        stdin  = MyPipe()
        stdout = MyPipe()

        # Initialize qiling instance with library path as given in the 
        # examples directory of the project. This is where the custom libraries
        # are provided which will aid emulation
        ql = Qiling([elf_path], "rootfs/x8664_linux", stdin=stdin, output="off", stdout=stdout)
        # Hook compare
        ql.hook_address(h_c, 0x5555555549b3)
        # Hook length
        ql.hook_address(h_d, 0x55555555496b)

        # Write test flag in input
        stdin.write(str.encode(test_flag))
        # Run emulation
        ql.run()

        # Check output
        out = str(stdout.read(1024))
        if "good job" in out:
            # We found what we wanted, now quit
            return test_flag

        # Clear current instances
        del stdin
        del stdout
        del ql

        # Failed output
        return None

    while True:
        # Run the emulation once
        sol = run_one_round(next_flag)
        if sol is not None:
            # WE GOT FLAG
            return sol

        # Next flag to try
        next_flag = ''.join(base_flag)
        base_flag = []

print(solve_one(base_path + "chal_0"))