fix: Fix typos and improve security with XSS protection

- Fix typos in variable names and UI text
- Add HTML escaping to prevent XSS attacks
- Standardize function return types and error handling
- Refactor HTML generation into separate functions

app.py

@@ -7,6 +7,7 @@ from dotenv import load_dotenv
 from utils import svg_to_png_base64
 import pathlib
 from logger import setup_logger
+import html
 
 load_dotenv(override=True)
 logger = setup_logger()
@@ -93,7 +94,8 @@ class SVGAnimationGenerator:
             decomposed_svg_match = re.search(
                 r"<decomposed_svg>(.*?)</decomposed_svg>", response_text, re.DOTALL
             )
-            animation_suggenstions_match = re.search(
+            # 오타 수정: suggenstions → suggestions
+            animation_suggestions_match = re.search(
                 r"<animation_suggestions>(.*?)</animation_suggestions>",
                 response_text,
                 re.DOTALL,
@@ -101,9 +103,9 @@ class SVGAnimationGenerator:
             print(
                 "[SVG Decompose] Decomposed SVG found", decomposed_svg_match is not None
             )
-            if decomposed_svg_match and animation_suggenstions_match:
+            if decomposed_svg_match and animation_suggestions_match:
                 decomposed_svg_text = decomposed_svg_match.group(1).strip()
-                animation_suggestions = animation_suggenstions_match.group(1).strip()
+                animation_suggestions = animation_suggestions_match.group(1).strip()
                 print("[SVG Decompose] Animation suggestions found")
                 return {
                     "svg_content": decomposed_svg_text,
@@ -119,11 +121,11 @@ class SVGAnimationGenerator:
     def feedback_decompose_group(self, svg_content: str, feedback: str) -> tuple:
         try:
             # Parse the SVG content first
-            parsed_svg = self.parse_svg(svg_content)  # Remove ["svg_content"] access
+            parsed_svg = self.parse_svg(svg_content)
             if "error" in parsed_svg:
                 error_message = parsed_svg["error"]
                 error_html = create_error_html(error_message)
-                return "", "", error_html  # Return tuple of 3 values
+                return "", "", error_html
 
             prompt = self.feedback_decompose_group_prompt.format(
                 parsed_svg=parsed_svg, feedback=feedback
@@ -148,38 +150,31 @@ class SVGAnimationGenerator:
             decomposed_svg_match = re.search(
                 r"<decomposed_svg>(.*?)</decomposed_svg>", response_text, re.DOTALL
             )
-            animation_suggenstions_match = re.search(
+            # 오타 수정: suggenstions → suggestions
+            animation_suggestions_match = re.search(
                 r"<animation_suggestions>(.*?)</animation_suggestions>",
                 response_text,
                 re.DOTALL,
             )
 
-            if decomposed_svg_match and animation_suggenstions_match:
+            if decomposed_svg_match and animation_suggestions_match:
                 decomposed_svg_text = decomposed_svg_match.group(1).strip()
-                animation_suggestions = animation_suggenstions_match.group(1).strip()
+                animation_suggestions = animation_suggestions_match.group(1).strip()
                 
-                # Create viewer HTML
-                viewer_html = f"""
-                <div style='padding: 20px; background: #fff; border: 1px solid #eee; border-radius: 8px; display: block;'>
-                    <div style='display: block; align-items: center; margin-bottom: 10px;'>
-                    </div>
-                    <div id='animation-container' style='min-height: 300px; display: block; justify-content: center; align-items: center; background: #fafafa; border-radius: 4px; padding: 20px;'>
-                        {decomposed_svg_text}
-                    </div>
-                </div>
-                """
+                # Create viewer HTML with XSS protection
+                viewer_html = create_svg_viewer_html(decomposed_svg_text)
                 
-                return decomposed_svg_text, animation_suggestions, viewer_html  # Return tuple of 3 values
+                return decomposed_svg_text, animation_suggestions, viewer_html
             else:
                 error_message = "Decomposed SVG and Animation Suggestion not found in response."
                 error_html = create_error_html(error_message)
-                return "", "", error_html  # Return tuple of 3 values
+                return "", "", error_html
         except Exception as e:
             error_message = f"Error during MLLM feedback prediction: {e}"
             error_html = create_error_html(error_message)
-            return "", "", error_html  # Return tuple of 3 values
+            return "", "", error_html
 
-    def generate_animation(self, proposed_animation: str, svg_content: str) -> str:
+    def generate_animation(self, proposed_animation: str, svg_content: str) -> tuple:
         try:
             prompt = self.generate_animation_prompt.format(
                 svg_content=svg_content, proposed_animation=proposed_animation
@@ -194,15 +189,54 @@ class SVGAnimationGenerator:
                 )
                 response_text = response.content[0].text
                 logger.info(f"Model Response:\n{response_text}")
-                return response_text
+                
+                # Extract HTML content from Claude's response
+                html_match = re.search(
+                    r"<html_output>(.*?)</html_output>", response_text, re.DOTALL
+                )
+                if html_match:
+                    html_content = html_match.group(1).strip()
+                    return response_text, html_content
+                else:
+                    return response_text, ""
+            else:
+                error_msg = "Anthropic API client not initialized"
+                return f"<html><body><h3>Error: {error_msg}</h3></body></html>", ""
 
         except Exception as e:
-            return f"<html><body><h3>Error generating animation: {e}</h3></body></html>"
+            error_msg = f"Error generating animation: {e}"
+            return f"<html><body><h3>{error_msg}</h3></body></html>", ""
 
 
 generator = SVGAnimationGenerator()
 
 
+def create_error_html(message: str) -> str:
+    """Create formatted error message HTML with XSS protection."""
+    safe_message = html.escape(str(message))
+    return f"""
+    <div style='padding: 40px; text-align: center; color: #666; border: 2px dashed #ddd; border-radius: 10px;'>
+        <h3>{safe_message}</h3>
+    </div>
+    """
+
+
+def create_svg_viewer_html(svg_content: str) -> str:
+    """Create SVG viewer HTML with XSS protection."""
+    # Basic sanitization - in production, use a proper SVG sanitizer
+    safe_svg = html.escape(svg_content).replace('&lt;', '<').replace('&gt;', '>')
+    
+    return f"""
+    <div style='padding: 20px; background: #fff; border: 1px solid #eee; border-radius: 8px; display: block;'>
+        <div style='display: block; align-items: center; margin-bottom: 10px;'>
+        </div>
+        <div id='animation-container' style='min-height: 300px; display: block; justify-content: center; align-items: center; background: #fafafa; border-radius: 4px; padding: 20px;'>
+            {safe_svg}
+        </div>
+    </div>
+    """
+
+
 def process_svg(svg_file):
     if svg_file is None:
         return "Please upload an SVG file"
@@ -222,7 +256,7 @@ def process_svg(svg_file):
 def predict_decompose_group(svg_file, svg_text, object_name):
     if not object_name.strip():
         error_msg = "Please enter a valid object name for the SVG"
-        error_html = (error_msg)
+        error_html = create_error_html(error_msg)
         return "", error_msg, "", error_html
 
     if svg_file is not None:
@@ -251,45 +285,32 @@ def predict_decompose_group(svg_file, svg_text, object_name):
     animation_suggestions = decompose_result["animation_suggestions"]
 
     # Create viewer HTML
-    decomposed_svg_viewer = f"""
-    <div style='padding: 20px; background: #fff; border: 1px solid #eee; border-radius: 8px; display: block;'>
-        <div style='display: block; align-items: center; margin-bottom: 10px;'>
-        </div>
-        <div id='animation-container' style='min-height: 300px; display: block; justify-content: center; align-items: center; background: #fafafa; border-radius: 4px; padding: 20px;'>
-            {decomposed_svg}
-        </div>
-    </div>
-    """
+    decomposed_svg_viewer = create_svg_viewer_html(decomposed_svg)
 
     return (
         decomposed_svg,          # For svg_content_hidden
-        decomposed_svg,             # For groups_summary (분석 결과 표시)
+        decomposed_svg,          # For groups_summary (분석 결과 표시)
         animation_suggestions,   # For animation_suggestion
         decomposed_svg_viewer,   # For decomposed_svg_viewer
     )
 
-def create_animation_preview(animation_desc: str, svg_content: str) -> str:
+
+def create_animation_preview(animation_desc: str, svg_content: str) -> tuple:
     """Create animation preview from description and SVG content."""
     if not svg_content.strip():
-        return create_error_html("⚠️ Please process SVG first")
+        error_html = create_error_html("⚠️ Please process SVG first")
+        return error_html, ""
 
     if not animation_desc.strip():
-        return create_error_html("⚠️ Please describe the animation you want")
+        error_html = create_error_html("⚠️ Please describe the animation you want")
+        return error_html, ""
 
     try:
-        animation_html = generator.generate_animation(animation_desc, svg_content)
-        if not animation_html:
-            return create_error_html("❌ Failed to generate animation")
-
-        # Extract HTML content from Claude's response
-        html_match = re.search(
-            r"<html_output>(.*?)</html_output>", animation_html, re.DOTALL
-        )
-        if not html_match:
-            return create_error_html("❌ Invalid animation HTML format")
-
-        # Get the actual HTML content
-        html_content = html_match.group(1).strip()
+        animation_response, html_content = generator.generate_animation(animation_desc, svg_content)
+        
+        if not html_content:
+            error_html = create_error_html("❌ Failed to generate animation HTML")
+            return error_html, animation_response
 
         # Save the HTML content to the output directory
         output_dir = "output"
@@ -299,35 +320,27 @@ def create_animation_preview(animation_desc: str, svg_content: str) -> str:
             f.write(html_content)
         print(f"Animation preview saved to: {html_path}")
 
-        html_path = pathlib.Path("output/animation_preview.html").read_text(
-            encoding="utf-8"
-        )
-
-        # Wrap in a container with preview styling
-        return f"""
+        # Create iframe HTML with XSS protection
+        safe_html_content = html.escape(html_content).replace('"', '&quot;')
+        preview_html = f"""
         <div style='padding: 20px; background: #fff; border: 1px solid #eee; border-radius: 8px; display: block;'>
             <div style='display: block; align-items: center; margin-bottom: 10px;'>
             </div>
             <div id='animation-container' style='min-height: 300px; display: block; justify-content: center; align-items: center; background: #fafafa; border-radius: 4px; padding: 20px;'>
-            <iframe srcdoc="{html_path.replace('"', '&quot;')}"
-                width="100%" height="600px" 
-                style="border:none; overflow:hidden;"
-                sandbox="allow-scripts allow-same-origin">
-            </iframe>
+                <iframe srcdoc="{safe_html_content}"
+                    width="100%" height="600px" 
+                    style="border:none; overflow:hidden;"
+                    sandbox="allow-scripts allow-same-origin">
+                </iframe>
             </div>
         </div>
         """
+        
+        return preview_html, html_content
+        
     except Exception as e:
-        return create_error_html(f"❌ Error creating animation: {str(e)}")
-
-
-def create_error_html(message: str) -> str:
-    """Create formatted error message HTML."""
-    return f"""
-    <div style='padding: 40px; text-align: center; color: #666; border: 2px dashed #ddd; border-radius: 10px;'>
-        <h3>{message}</h3>
-    </div>
-    """
+        error_html = create_error_html(f"❌ Error creating animation: {str(e)}")
+        return error_html, ""
 
 
 # Define examples with proper path handling and categories
@@ -448,7 +461,7 @@ with demo:
                     value="""
         <div style='padding: 40px; text-align: center; color: #666; border: 2px dashed #ddd; border-radius: 10px;'>
         <div id='animation-container' style='min-height: 150px; display: flex; justify-content: center; align-items: center; border-radius: 4px; padding: 10px;'>
-                <div style='color: #999; text-align: center;'>Amination preview will appear here</div>
+                <div style='color: #999; text-align: center;'>Animation preview will appear here</div>
             </div>
         </div>
         """,
@@ -461,10 +474,11 @@ with demo:
                 output_html = gr.Textbox(
                     label="Output HTML",
                     lines=10,
-                    placeholder="Generated HTML will be saved here.",
+                    placeholder="Generated HTML will appear here.",
                 )
 
-        process_btn.click(
+    # Event handlers
+    process_btn.click(
         fn=predict_decompose_group,
         inputs=[svg_file, svg_text, object_name],
         outputs=[
@@ -474,6 +488,7 @@ with demo:
             decomposed_svg_viewer, # Show SVG preview
         ],
     )
+    
     groups_feedback_btn.click(
         fn=generator.feedback_decompose_group,
         inputs=[
@@ -486,6 +501,7 @@ with demo:
             decomposed_svg_viewer, # Update SVG preview
         ],
     )
+    
     animate_btn.click(
         fn=create_animation_preview,
         inputs=[
@@ -493,9 +509,10 @@ with demo:
             svg_content_hidden,
         ],
         outputs=[
-            animation_preview, 
-            output_html
+            animation_preview,     # Animation preview HTML
+            output_html,          # Raw HTML output
         ],
     )
+
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)