1+ import io
12from typing import TYPE_CHECKING , ClassVar , Dict , List , Optional , Type
23
34import typer
5+ from rich .console import Console
46
7+ from cycode .cli import consts
8+ from cycode .cli .cli_types import ExportTypeOption
9+ from cycode .cli .console import console , console_err
510from cycode .cli .exceptions .custom_exceptions import CycodeError
611from cycode .cli .models import CliError , CliResult
712from cycode .cli .printers .json_printer import JsonPrinter
@@ -27,57 +32,115 @@ class ConsolePrinter:
2732 'rich_sca' : ScaTablePrinter ,
2833 }
2934
30- def __init__ (self , ctx : typer .Context ) -> None :
35+ def __init__ (
36+ self ,
37+ ctx : typer .Context ,
38+ console_override : Optional ['Console' ] = None ,
39+ console_err_override : Optional ['Console' ] = None ,
40+ output_type_override : Optional [str ] = None ,
41+ ) -> None :
3142 self .ctx = ctx
43+ self .console = console_override or console
44+ self .console_err = console_err_override or console_err
3245
3346 self .scan_type = self .ctx .obj .get ('scan_type' )
34- self .output_type = self .ctx .obj .get ('output' )
47+ self .output_type = output_type_override or self .ctx .obj .get ('output' )
3548 self .aggregation_report_url = self .ctx .obj .get ('aggregation_report_url' )
3649
37- self ._printer_class = self ._AVAILABLE_PRINTERS .get (self .output_type )
38- if self ._printer_class is None :
39- raise CycodeError (f'"{ self .output_type } )
50+ self .printer = self ._get_scan_printer ()
4051
41- def print_scan_results (
42- self ,
43- local_scan_results : List ['LocalScanResult' ],
44- errors : Optional [Dict [str , 'CliError' ]] = None ,
45- ) -> None :
46- printer = self ._get_scan_printer ()
47- printer .print_scan_results (local_scan_results , errors )
52+ self .console_record = None
53+
54+ self .export_type = self .ctx .obj .get ('export_type' )
55+ self .export_file = self .ctx .obj .get ('export_file' )
56+ if console_override is None and self .export_type and self .export_file :
57+ self .console_record = ConsolePrinter (
58+ ctx ,
59+ console_override = Console (record = True , file = io .StringIO ()),
60+ console_err_override = Console (stderr = True , record = True , file = io .StringIO ()),
61+ output_type_override = 'json' if self .export_type == 'json' else self .output_type ,
62+ )
4863
4964 def _get_scan_printer (self ) -> 'PrinterBase' :
50- printer_class = self ._printer_class
65+ printer_class = self ._AVAILABLE_PRINTERS . get ( self . output_type )
5166
5267 composite_printer = self ._AVAILABLE_PRINTERS .get (f'{ self .output_type } { self .scan_type } )
5368 if composite_printer :
5469 printer_class = composite_printer
5570
56- return printer_class (self .ctx )
71+ if not printer_class :
72+ raise CycodeError (f'"{ self .output_type } )
73+
74+ return printer_class (self .ctx , self .console , self .console_err )
75+
76+ def print_scan_results (
77+ self ,
78+ local_scan_results : List ['LocalScanResult' ],
79+ errors : Optional [Dict [str , 'CliError' ]] = None ,
80+ ) -> None :
81+ if self .console_record :
82+ self .console_record .print_scan_results (local_scan_results , errors )
83+ self .printer .print_scan_results (local_scan_results , errors )
5784
5885 def print_result (self , result : CliResult ) -> None :
59- self ._printer_class (self .ctx ).print_result (result )
86+ if self .console_record :
87+ self .console_record .print_result (result )
88+ self .printer .print_result (result )
6089
6190 def print_error (self , error : CliError ) -> None :
62- self ._printer_class (self .ctx ).print_error (error )
91+ if self .console_record :
92+ self .console_record .print_error (error )
93+ self .printer .print_error (error )
6394
6495 def print_exception (self , e : Optional [BaseException ] = None , force_print : bool = False ) -> None :
6596 """Print traceback message in stderr if verbose mode is set."""
6697 if force_print or self .ctx .obj .get ('verbose' , False ):
67- self ._printer_class (self .ctx ).print_exception (e )
98+ if self .console_record :
99+ self .console_record .print_exception (e )
100+ self .printer .print_exception (e )
101+
102+ def export (self ) -> None :
103+ if self .console_record is None :
104+ raise CycodeError ('Console recording was not enabled. Cannot export.' )
105+
106+ if not self .export_file .suffix :
107+ # resolve file extension based on the export type if not provided in the file name
108+ self .export_file = self .export_file .with_suffix (f'.{ self .export_type .lower ()} )
109+
110+ if self .export_type is ExportTypeOption .HTML :
111+ self .console_record .console .save_html (self .export_file )
112+ elif self .export_type is ExportTypeOption .SVG :
113+ self .console_record .console .save_svg (self .export_file , title = consts .APP_NAME )
114+ elif self .export_type is ExportTypeOption .JSON :
115+ with open (self .export_file , 'w' , encoding = 'UTF-8' ) as f :
116+ self .console_record .console .file .seek (0 )
117+ f .write (self .console_record .console .file .read ())
118+ else :
119+ raise CycodeError (f'Export type "{ self .export_type } )
120+
121+ export_format_msg = f'{ self .export_type .upper ()}
122+ if self .export_type in {ExportTypeOption .HTML , ExportTypeOption .SVG }:
123+ export_format_msg += f' with { self .output_type .upper ()}
124+
125+ clickable_path = f'[link=file://{ self .export_file } { self .export_file }
126+ self .console .print (f'[b green]Cycode CLI output exported to { clickable_path } { export_format_msg } )
127+
128+ @property
129+ def is_recording (self ) -> bool :
130+ return self .console_record is not None
68131
69132 @property
70133 def is_json_printer (self ) -> bool :
71- return self ._printer_class == JsonPrinter
134+ return isinstance ( self .printer , JsonPrinter )
72135
73136 @property
74137 def is_table_printer (self ) -> bool :
75- return self ._printer_class == TablePrinter
138+ return isinstance ( self .printer , TablePrinter )
76139
77140 @property
78141 def is_text_printer (self ) -> bool :
79- return self ._printer_class == TextPrinter
142+ return isinstance ( self .printer , TextPrinter )
80143
81144 @property
82145 def is_rich_printer (self ) -> bool :
83- return self ._printer_class == RichPrinter
146+ return isinstance ( self .printer , RichPrinter )
0 commit comments