utils/report_generator.py

from huggingface_hub import InferenceClient
import logging
from typing import Dict, List
import os


class ReportGenerator:

    def __init__(self):
        try:
            hf_token = os.getenv("HF_API_TOKEN")

            if not hf_token:
                raise ValueError("Hugging Face API token is required")

            self.client = InferenceClient(
                model="mistralai/Mistral-7B-Instruct-v0.3", token=hf_token
            )
        except Exception as e:
            logging.error(f"Failed to initialize Inference client: {str(e)}")
            raise

    def _truncate_list(self, items: List[str], max_items: int = 15) -> List[str]:
        """Truncate a list to a maximum number of items."""
        if len(items) > max_items:
            return items[:max_items] + [f"... and {len(items) - max_items} more"]
        return items

    def generate_report(self, basic_info: Dict, security_issues: List[str]) -> str:
        try:
            truncated_permissions = self._truncate_list(
                basic_info.get("permissions", [])
            )
            truncated_issues = self._truncate_list(security_issues)

            prompt = f"""<s>[INST] As a security expert and Android developer, analyze this Android app and provide specific code-level recommendations:

App: {basic_info.get('app_name', 'Unknown')} ({basic_info.get('package', 'Unknown')})
Version: {basic_info.get('version', 'Unknown')}

Security Issues Found:
{chr(10).join([f"- {issue}" for issue in truncated_issues])}

Requested Permissions:
{chr(10).join([f"- {perm}" for perm in truncated_permissions])}

Please provide:
1. A brief security risk assessment
2. Specific code-level fixes for each issue, including example code snippets where applicable
3. Best practices that should be implemented
4. Recommended security configurations

Format your response with clear sections and include Android/Java code examples for fixes. [/INST]</s>"""

            # Call the Inference API with adjusted parameters for more detailed output
            response = self.client.text_generation(
                prompt,
                max_new_tokens=1024,  # Increased token limit for more detailed response
                temperature=0.2,  # Slightly increased for more creative suggestions
                repetition_penalty=1.2,
                do_sample=True,
                return_full_text=False,
            )

            if not response:
                return "Error: Could not generate a meaningful report"

            # Format the final report with clear sections
            return f"""## Security Analysis Report

{response}

### Additional Resources
- [Android Security Best Practices](https://developer.android.com/topic/security/best-practices)
- [OWASP Mobile Security Testing Guide](https://owasp.org/www-project-mobile-security-testing-guide/)

---
*Note: This analysis is based on static analysis and may not cover all runtime behaviors. Implement these suggestions after thorough testing.*"""

        except Exception as e:
            logging.error(f"Error generating report: {str(e)}")
            return f"Error generating report: {str(e)}"