#!/usr/bin/env python3
#-*- coding: utf-8 -*-

import sys
import os

from qiling import *
from capstone import *

base_address = 0x400000

offset = 0x402176 - 0x402044

address_of_call = 0x4012a7

flag = ""
pos = 0

def my_fgets(ql):
    '''
    Write a valid size buffer (with incorrect content),
    then set RAX to correct value, this will make that
    program continues until the checks.
    '''
    buff_addr = ql.os.function_arg[0]
    print("[+] Buffer address to write string: 0x%lx" % (buff_addr))
    print("[+] Writing %s" % ('A'*0x26))
    ql.mem.string(buff_addr, 'A'*0x26)
    ql.reg.write('RAX', buff_addr)

def patch2(ql):
    '''
    Patch to extract data from the comparison
    '''
    global flag
    global pos

    EDX = ql.reg.read("EDX")
    EAX = ql.reg.read("EAX")
    print("[+] Value that should be written in position %d: %x (%c)" % (pos, EDX, chr(EDX)))
    ql.reg.write("EAX", EDX)
    flag += chr(EDX)
    pos += 1

def patch1(ql):
    '''
    Patch to extract address to call
    '''
    RCX = ql.reg.read("RCX")
    print("[+] Address with decryption function: 0x%lx" % (RCX))
    print("[+] New address to patch: 0x%lx" % (RCX + offset))
    buf = ql.mem.read((RCX + offset), 2)
    md = Cs(CS_ARCH_X86, CS_MODE_64)
    for i in md.disasm(buf, (RCX + offset)):
        print("[+] Disassembly of patched address: :: 0x%x:\t%s\t%s" %(i.address, i.mnemonic, i.op_str))
    ql.hook_address(callback=patch2, address=(RCX + offset))


ql = Qiling(filename=["./x-and-or"], archtype="x86-64", ostype="linux", rootfs='/', profile='my_profiler.ql')
ql.hook_address(callback=patch1, address=address_of_call)
ql.set_api('fgets', my_fgets)
ql.run()

print("Flag is: '%s'" % (flag))