report_gen.py

import io
from reportlab.lib import colors
from reportlab.lib.pagesizes import A4
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Image, Table, TableStyle
from reportlab.lib.units import mm
from reportlab.graphics.shapes import Drawing, Rect, String, Line

class SliderFlowable(Drawing):
    def __init__(self, name, value, min_val=1, max_val=10, width=400, height=50):
        Drawing.__init__(self, width, height)
        self.name = name
        self.value = value
        self.min_val = min_val
        self.max_val = max_val
        
        # Draw slider bar
        self.add(Rect(50, 15, 300, 20, fillColor=colors.lightgrey, strokeColor=colors.grey))
        
        # Draw slider value
        value_width = 50 + ((self.value - self.min_val) / (self.max_val - self.min_val) * 300)
        self.add(Rect(50, 15, value_width - 50, 20, fillColor=colors.blue, strokeColor=None))
        
        # Add slider name and value
        self.add(String(0, 40, self.name, fontSize=10))
        self.add(String(value_width - 10, 40, str(self.value), fontSize=10))

def create_styles():
    styles = getSampleStyleSheet()
    styles.add(ParagraphStyle(
        name='SectionTitle',
        parent=styles['Heading1'],
        fontSize=16,
        textColor=colors.HexColor("#3b0b75"),
        spaceAfter=10
    ))
    styles.add(ParagraphStyle(
        name='SubsectionTitle',
        parent=styles['Heading2'],
        fontSize=14,
        textColor=colors.HexColor("#52176a"),
        spaceAfter=8
    ))
    styles.add(ParagraphStyle(
        name='BodyText',
        parent=styles['Normal'],
        fontSize=10,
        textColor=colors.HexColor("#26004d"),
        spaceAfter=6
    ))
    styles.add(ParagraphStyle(
        name='AnswerText',
        parent=styles['BodyText'],
        backColor=colors.HexColor("#f0f2fd"),
        borderPadding=(4, 4, 4, 4),
        spaceAfter=10
    ))
    return styles

def generate_pdf(pages, answers):
    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=A4, rightMargin=20*mm, leftMargin=20*mm, topMargin=20*mm, bottomMargin=20*mm)
    styles = create_styles()
    story = [Paragraph("Critical Infrastructure - LLM Considerations Report", styles['Title'])]

    for page in pages[:-1]:  # Skip the last page (report generation page)
        if page['type'] == 'description':
            story.append(Paragraph(page['title'], styles['SectionTitle']))
            story.append(Paragraph(page['content'], styles['BodyText']))
        
        elif page['type'] == 'quiz':
            story.append(Paragraph(page['title'], styles['SubsectionTitle']))
            story.append(Paragraph(page['description'], styles['BodyText']))
            
            answer = answers.get(page['input_key'], "")
            story.append(Paragraph(f"<b>Analysis:</b>", styles['BodyText']))
            story.append(Paragraph(answer, styles['AnswerText']))
            
            risk_level = answers.get(page['slider_key'], 1)
            story.append(SliderFlowable(f"Risk Level: {risk_level}", risk_level))
        
        elif page['type'] == 'inference':
            story.append(Paragraph(page['title'], styles['SubsectionTitle']))
            story.append(Paragraph(page['prompt_description'], styles['BodyText']))
            
            user_input = answers.get(page['input_key'], "")
            inference_output = answers.get(page['output_key'], "")
            
            story.append(Paragraph("<b>User Input:</b>", styles['BodyText']))
            story.append(Paragraph(user_input, styles['AnswerText']))
            story.append(Paragraph("<b>Inference Output:</b>", styles['BodyText']))
            story.append(Paragraph(inference_output, styles['AnswerText']))
        
        elif page['type'] == 'architecture':
            story.append(Paragraph(page['title'], styles['SubsectionTitle']))
            
            # We can't include the actual architecture diagram in the PDF,
            # so we'll just mention that it was reviewed
            story.append(Paragraph("Architecture diagram was reviewed.", styles['BodyText']))
            
            analysis = answers.get(page['input_key'], "")
            story.append(Paragraph("<b>Security Analysis:</b>", styles['BodyText']))
            story.append(Paragraph(analysis, styles['AnswerText']))
        
        story.append(Spacer(1, 12))

    doc.build(story)
    buffer.seek(0)
    return buffer

def create_page_template(canvas, doc):
    canvas.saveState()
    canvas.setFont('Helvetica', 9)
    canvas.drawString(20*mm, 10*mm, f"Page {doc.page}")
    canvas.restoreState()