Coverage for src/alprina_cli/fix_command.py: 16%

1""" 

2Fix command - AI-powered vulnerability remediation. 

3Generates and applies secure code fixes for vulnerabilities. 

4""" 

5 

6from pathlib import Path 

7from typing import Optional, Dict, List 

8from rich.console import Console 

9from rich.panel import Panel 

10from rich.table import Table 

11from rich.syntax import Syntax 

12from rich.prompt import Confirm 

13import json 

14 

15from .services.fix_generator import get_fix_generator 

16from loguru import logger 

17 

18console = Console() 

19 

20 

21def fix_command( 

22 target: str, 

23 finding_id: Optional[str] = None, 

24 auto_fix: bool = False, 

25 severity: Optional[str] = None, 

26 preview_only: bool = False 

27): 

28 """ 

29 Generate AI-powered fixes for vulnerabilities. 

30 

31 Args: 

32 target: Path to scan or specific file 

33 finding_id: Specific finding ID to fix 

34 auto_fix: Automatically apply fixes without confirmation 

35 severity: Fix only specific severity (critical, high, medium, low) 

36 preview_only: Show fixes without applying 

37 """ 

38 console.print(Panel( 

39 f"🤖 AI-Powered Fix Generator\n\n" 

40 f"Target: [bold]{target}[/bold]\n" 

41 f"Mode: {'[green]Auto-apply[/green]' if auto_fix else '[yellow]Interactive[/yellow]'}", 

42 title="Alprina Fix", 

43 border_style="cyan" 

44 )) 

45 

46 try: 

47 # TODO: Integrate with existing scan results 

48 # For now, demonstrate with sample finding 

49 _demo_fix_generation(target, auto_fix, preview_only) 

50 

51 except Exception as e: 

52 console.print(f"[red]Fix generation failed: {e}[/red]") 

53 logger.error(f"Fix command error: {e}", exc_info=True) 

54 

55 

56def suggest_fixes_command(scan_results_file: str): 

57 """ 

58 Suggest fixes for findings in a scan results file. 

59 

60 Args: 

61 scan_results_file: Path to JSON file with scan results 

62 """ 

63 try: 

64 # Load scan results 

65 with open(scan_results_file, 'r') as f: 

66 results = json.load(f) 

67 

68 findings = results.get("findings", []) 

69 

70 if not findings: 

71 console.print("[yellow]No findings in scan results[/yellow]") 

72 return 

73 

74 console.print(Panel( 

75 f"📋 Analyzing {len(findings)} findings\n" 

76 f"Generating AI-powered fix suggestions...", 

77 title="Fix Suggestions", 

78 border_style="cyan" 

79 )) 

80 

81 # Generate fixes for all findings 

82 fix_generator = get_fix_generator() 

83 

84 for i, finding in enumerate(findings, 1): 

85 console.print(f"\n[bold cyan]Finding {i}/{len(findings)}:[/bold cyan]") 

86 _process_single_finding(finding, fix_generator, preview_only=True) 

87 

88 except FileNotFoundError: 

89 console.print(f"[red]Scan results file not found: {scan_results_file}[/red]") 

90 except json.JSONDecodeError: 

91 console.print(f"[red]Invalid JSON in scan results file[/red]") 

92 except Exception as e: 

93 console.print(f"[red]Error: {e}[/red]") 

94 logger.error(f"Suggest fixes error: {e}", exc_info=True) 

95 

96 

97def _demo_fix_generation(target: str, auto_fix: bool, preview_only: bool): 

98 """ 

99 Demo fix generation with sample vulnerability. 

100 TODO: Replace with real scan integration. 

101 """ 

102 # Sample vulnerability for demonstration 

103 sample_finding = { 

104 "type": "SQL Injection", 

105 "severity": "CRITICAL", 

106 "description": "User input directly concatenated into SQL query without sanitization", 

107 "location": f"{target}:45", 

108 "line": 45, 

109 "cwe": "CWE-89", 

110 "cvss_score": 9.8 

111 } 

112 

113 # Sample vulnerable code 

114 sample_code = '''import sqlite3 

115 

116def get_user_by_id(user_id): 

117 """Get user from database by ID.""" 

118 conn = sqlite3.connect('users.db') 

119 cursor = conn.cursor() 

120  

121 # VULNERABLE: Direct string formatting 

122 query = f"SELECT * FROM users WHERE id = {user_id}" 

123  

124 cursor.execute(query) 

125 result = cursor.fetchone() 

126 conn.close() 

127 return result 

128''' 

129 

130 fix_generator = get_fix_generator() 

131 

132 console.print("\n[bold]🔍 Analyzing Vulnerability...[/bold]") 

133 console.print(f" Type: [red]{sample_finding['type']}[/red]") 

134 console.print(f" Severity: [red]{sample_finding['severity']}[/red]") 

135 console.print(f" Location: [cyan]{sample_finding['location']}[/cyan]") 

136 if sample_finding.get("cwe"): 

137 console.print(f" CWE: [cyan]{sample_finding['cwe']}[/cyan]") 

138 if sample_finding.get("cvss_score"): 

139 console.print(f" CVSS: [red]{sample_finding['cvss_score']}/10.0[/red]") 

140 

141 console.print("\n[bold]🤖 Generating AI-Powered Fix...[/bold]") 

142 

143 # Generate fix 

144 fix_data = fix_generator.generate_fix( 

145 code=sample_code, 

146 vulnerability=sample_finding, 

147 filename=target 

148 ) 

149 

150 # Display results 

151 _display_fix(sample_code, fix_data, auto_fix, preview_only) 

152 

153 

154def _process_single_finding(finding: Dict, fix_generator, preview_only: bool = True): 

155 """Process a single finding and generate fix.""" 

156 console.print(f" Type: [yellow]{finding.get('type')}[/yellow]") 

157 console.print(f" Severity: [red]{finding.get('severity')}[/red]") 

158 console.print(f" Location: [cyan]{finding.get('location')}[/cyan]") 

159 

160 # TODO: Load actual file content 

161 # For now, skip file reading 

162 console.print(" [dim]Fix generation requires file access (not implemented in this demo)[/dim]") 

163 

164 

165def _display_fix(original_code: str, fix_data: dict, auto_fix: bool, preview_only: bool): 

166 """Display fix with before/after comparison.""" 

167 

168 if fix_data.get("error"): 

169 console.print(f"\n[red]❌ Error generating fix:[/red] {fix_data['error']}") 

170 return 

171 

172 confidence = fix_data.get("confidence", 0.0) 

173 

174 # Show confidence score 

175 confidence_color = "green" if confidence >= 0.8 else "yellow" if confidence >= 0.6 else "red" 

176 console.print(f"\n[bold]✨ Fix Generated (Confidence: [{confidence_color}]{confidence*100:.0f}%[/{confidence_color}])[/bold]") 

177 

178 # Show explanation 

179 if fix_data.get("explanation"): 

180 console.print(f"\n[bold cyan]💡 Why This Fix Works:[/bold cyan]") 

181 console.print(Panel( 

182 fix_data["explanation"], 

183 border_style="cyan", 

184 padding=(1, 2) 

185 )) 

186 

187 # Show changes 

188 if fix_data.get("changes"): 

189 console.print(f"\n[bold cyan]📝 Changes Made:[/bold cyan]") 

190 for i, change in enumerate(fix_data["changes"], 1): 

191 console.print(f" {i}. {change}") 

192 

193 # Show before/after code 

194 console.print(f"\n[bold red]❌ Before (Vulnerable):[/bold red]") 

195 syntax = Syntax(original_code.strip(), "python", theme="monokai", line_numbers=True) 

196 console.print(Panel(syntax, border_style="red")) 

197 

198 fixed_code = fix_data.get("fixed_code", "") 

199 if fixed_code: 

200 console.print(f"\n[bold green]✅ After (Secure):[/bold green]") 

201 syntax = Syntax(fixed_code.strip(), "python", theme="monokai", line_numbers=True) 

202 console.print(Panel(syntax, border_style="green")) 

203 

204 # Show diff 

205 if fix_data.get("diff"): 

206 console.print(f"\n[bold cyan]🔄 Diff:[/bold cyan]") 

207 console.print(Panel(fix_data["diff"], border_style="cyan")) 

208 

209 # Show security notes 

210 if fix_data.get("security_notes"): 

211 console.print(f"\n[bold yellow]⚠️ Security Notes:[/bold yellow]") 

212 for i, note in enumerate(fix_data["security_notes"], 1): 

213 console.print(f" {i}. {note}") 

214 

215 # Apply fix (interactive or auto) 

216 if not preview_only: 

217 _apply_fix_interactive(fix_data, auto_fix) 

218 

219 

220def _apply_fix_interactive(fix_data: dict, auto_fix: bool): 

221 """Interactively apply fix with user confirmation.""" 

222 

223 if auto_fix: 

224 console.print("\n[green]✓[/green] Auto-applying fix...") 

225 # TODO: Actually apply the fix 

226 console.print("[green]✓ Fix applied![/green]") 

227 console.print("[dim]💾 Backup saved as: file.py.backup[/dim]") 

228 else: 

229 console.print("\n[bold]Apply this fix?[/bold]") 

230 console.print(" [y] Yes, apply the fix") 

231 console.print(" [n] No, skip this fix") 

232 console.print(" [d] Show diff again") 

233 console.print(" [e] Explain in detail") 

234 

235 if Confirm.ask("Apply fix?", default=False): 

236 # TODO: Actually apply the fix 

237 console.print("[green]✓ Fix applied![/green]") 

238 console.print("[dim]💾 Backup saved as: file.py.backup[/dim]") 

239 else: 

240 console.print("[yellow]⊗ Fix skipped[/yellow]") 

241 

242 

243def _show_fix_summary(total: int, applied: int, skipped: int): 

244 """Show summary of fix session.""" 

245 console.print("\n" + "="*60) 

246 console.print("[bold]📊 Fix Summary[/bold]") 

247 console.print("="*60) 

248 

249 table = Table(show_header=False, box=None) 

250 table.add_row("Total findings:", str(total)) 

251 table.add_row("Fixes applied:", f"[green]{applied}[/green]") 

252 table.add_row("Fixes skipped:", f"[yellow]{skipped}[/yellow]") 

253 

254 console.print(table) 

255 

256 if applied > 0: 

257 console.print("\n[green]✓[/green] Run your tests to verify the fixes work correctly") 

258 console.print("[green]✓[/green] Backups saved with .backup extension") 

259 

260 console.print("\n[dim]Tip: Run 'alprina scan' again to verify fixes resolved the issues[/dim]")