import random
import string
import os
import traceback
from core.rwl import story_parser

# Configuration
NUM_FUZZ_CASES = 1000
MAX_STRING_LENGTH = 200
MALFORMED_PRD_DIR = '/mnt/e/genesis-system/core/rwl/attacks/malformed_prds/'
REPORT_FILE = '/mnt/e/genesis-system/core/rwl/attacks/parser_attack_report.md'


def generate_random_string(length):
    return ''.join(random.choice(string.printable) for _ in range(length))


def generate_malformed_prd():
    # Generate various types of malformed PRDs
    prd_type = random.choice(['random_string', 'empty', 'nested_tags', 'unbalanced_tags', 'invalid_tags'])

    if prd_type == 'random_string':
        return generate_random_string(random.randint(1, MAX_STRING_LENGTH))
    elif prd_type == 'empty':
        return ''
    elif prd_type == 'nested_tags':
        return '<story><nested><tags>{}</tags></nested></story>'.format(generate_random_string(random.randint(1, MAX_STRING_LENGTH)))
    elif prd_type == 'unbalanced_tags':
        return '<story><a><b></story>'
    elif prd_type == 'invalid_tags':
        return '<story><invalid_tag>{}</invalid_tag></story>'.format(generate_random_string(random.randint(1, MAX_STRING_LENGTH)))
    else:
        return '<story>{}</story>'.format(generate_random_string(random.randint(1, MAX_STRING_LENGTH)))


def fuzz_parser(num_cases):
    crash_cases = []
    edge_cases = []
    
    for i in range(num_cases):
        malformed_prd = generate_malformed_prd()
        try:
            story_parser.parse_story(malformed_prd)
        except Exception as e:
            crash_cases.append({
                'input': malformed_prd,
                'exception': str(e),
                'traceback': traceback.format_exc()
            })
        else:
            # Check for potential edge cases (e.g., parser handles unexpected input gracefully but produces odd results)
            # This requires deeper analysis of the parser's output.  For now, we'll just flag inputs that contain specific keywords.
            if 'invalid' in malformed_prd or 'unbalanced' in malformed_prd or 'nested' in malformed_prd:
                edge_cases.append({
                    'input': malformed_prd,
                    'note': 'Potential edge case - parser might misinterpret this input.'
                })

    return crash_cases, edge_cases


def save_malformed_prds(malformed_prds, directory):
    if not os.path.exists(directory):
        os.makedirs(directory)

    for i, prd in enumerate(malformed_prds):
        filename = os.path.join(directory, f'malformed_prd_{i}.xml')
        with open(filename, 'w') as f:
            f.write(prd)


def generate_report(crash_cases, edge_cases):
    report = "# Parser Fuzzing Attack Report\n\n"
    report += "## Summary\n\n"
    report += f"Number of fuzz cases: {NUM_FUZZ_CASES}\n"
    report += f"Number of crash cases: {len(crash_cases)}\n"
    report += f"Number of potential edge cases: {len(edge_cases)}\n\n"

    report += "## Crash Cases\n\n"
    if crash_cases:
        for i, case in enumerate(crash_cases):
            report += f"### Case {i + 1}\n\n"
            report += f"**Input:**\n\`\`\`\n{case['input']}\n\`\`\`\n\n"
            report += f"**Exception:** {case['exception']}\n\n"
            report += f"**Traceback:**\n\`\`\`\n{case['traceback']}\n\`\`\`\n\n"
    else:
        report += "No crash cases found.\n\n"

    report += "## Potential Edge Cases\n\n"
    if edge_cases:
        for i, case in enumerate(edge_cases):
            report += f"### Case {i + 1}\n\n"
            report += f"**Input:**\n\`\`\`\n{case['input']}\n\`\`\`\n\n"
            report += f"**Note:** {case['note']}\n\n"
    else:
        report += "No potential edge cases found.\n\n"

    return report


if __name__ == '__main__':
    crash_cases, edge_cases = fuzz_parser(NUM_FUZZ_CASES)
    #save_malformed_prds([case['input'] for case in crash_cases], MALFORMED_PRD_DIR)
    report = generate_report(crash_cases, edge_cases)

    with open(REPORT_FILE, 'w') as f:
        f.write(report)

    print(f"Fuzzing complete. Report saved to {REPORT_FILE}")