!55 SDK refactor POC - Findings Closes #103

Author: Vladimir Kontić

probely/__init__.py

@@ -10,11 +10,10 @@
     FindingSeverityEnum,
     FindingStateEnum,
     ScanStatusEnum,
+    TargetAPISchemaTypeEnum,
     TargetRiskEnum,
     TargetTypeEnum,
-    TargetAPISchemaTypeEnum,
 )
-from .sdk.findings import list_findings, retrieve_finding, retrieve_findings
 from .sdk.scans import (
     cancel_scan,
     cancel_scans,
@@ -61,9 +60,6 @@
     "resume_scans",
     "pause_scan",
     "pause_scans",
-    "list_findings",
-    "retrieve_finding",
-    "retrieve_findings",
     "ProbelyException",
     "ProbelyObjectNotFound",
     "ProbelyMissConfig",

probely/cli/commands/findings/get.py

@@ -6,7 +6,7 @@
 from probely.cli.renderers import OutputRenderer
 from probely.cli.tables.finding_table import FindingTable
 from probely.exceptions import ProbelyCLIValidation
-from probely.sdk.findings import list_findings, retrieve_findings
+from probely.sdk.managers import FindingManager
 
 
 def findings_get_command_handler(args: argparse.Namespace):
@@ -15,9 +15,9 @@ def findings_get_command_handler(args: argparse.Namespace):
         raise ProbelyCLIValidation("filters and Finding IDs are mutually exclusive.")
 
     if args.findings_ids:
-        findings_generator = retrieve_findings(findings_ids=args.findings_ids)
+        findings_generator = FindingManager().get_multiple(ids=args.findings_ids)
     else:
-        findings_generator = list_findings(findings_filters=filters)
+        findings_generator = FindingManager().list(filters=filters)
 
     output_type = OutputEnum[args.output] if args.output else None
     renderer = OutputRenderer(

probely/cli/commands/targets/get.py

@@ -1,7 +1,6 @@
 from probely.cli.commands.targets.schemas import TargetApiFiltersSchema
 from probely.cli.common import prepare_filters_for_api
 from probely.cli.enums import OutputEnum
-
 from probely.cli.renderers import OutputRenderer
 from probely.cli.tables.targets_table import TargetTable
 from probely.exceptions import ProbelyCLIValidation

probely/cli/renderers.py

@@ -1,6 +1,7 @@
 import json
 import sys
 import textwrap
+from datetime import datetime
 from typing import (
     Dict,
     Generator,
@@ -17,6 +18,7 @@
 from probely.cli.enums import OutputEnum
 from probely.cli.tables.base_table import BaseOutputTable
 from probely.sdk.enums import ProbelyCLIEnum
+from probely.sdk._schemas import Finding, FindingLabel
 
 UNKNOWN_VALUE_REP = "UNKNOWN"
 TARGET_NEVER_SCANNED_OUTPUT = "Never_scanned"
@@ -29,7 +31,7 @@ class OutputRenderer:
 
     def __init__(
         self,
-        records: Generator[dict, None, None],
+        records: Generator[Union[dict, Finding], None, None],
         output_type: Optional[OutputEnum],
         console: Console,
         table_cls: Type[BaseOutputTable],
@@ -53,12 +55,19 @@ def _render_json(self) -> None:
         for record in self.records:
             if not first:
                 self.console.print(",")
-            self.console.print(json.dumps(record, indent=2))
+
+            if hasattr(record, "to_json"):
+                # NOTE: just temporary solution while we finish SDK refactor and start using OOP approach everywhere
+                self.console.print(record.to_json(indent=2))
+            else:
+                self.console.print(json.dumps(record, indent=2))
             first = False
         self.console.print("]")
 
     def _render_yaml(self) -> None:
         for record in self.records:
+            if hasattr(record, "to_dict"):
+                record = record.to_dict(mode="json")
             self.console.print(yaml.dump([record], indent=2, width=sys.maxsize))
 
     def _render_table(self) -> None:
@@ -79,16 +88,18 @@ def get_printable_enum_value(enum: Type[ProbelyCLIEnum], api_enum_value: str) ->
         return UNKNOWN_VALUE_REP  # TODO: scenario that risk enum updated but CLI is forgotten
 
 
-def get_printable_labels(labels: List[Dict] = None) -> str:
+def get_printable_labels(labels: List[Union[Dict, FindingLabel]] = None) -> str:
     if labels is None:
         return "UNKNOWN_LABELS"
 
     labels_names = []
     try:
         for label in labels:
-            truncated_label = textwrap.shorten(
-                label["name"], width=16, placeholder="..."
-            )
+            if isinstance(label, FindingLabel):
+                label_name = label.name
+            else:
+                label_name = label["name"]
+            truncated_label = textwrap.shorten(label_name, width=16, placeholder="...")
             labels_names.append(truncated_label)
     except:
         return "UNKNOWN_LABELS"
@@ -99,12 +110,18 @@ def get_printable_labels(labels: List[Dict] = None) -> str:
 
 
 def get_printable_date(
-    date_string: Union[str, None],
+    date_input: Union[str, datetime, None],
     default_string: Union[str, None] = None,
 ) -> str:
-    if date_string:
-        datetime = parser.isoparse(date_string)
-        return datetime.strftime("%Y-%m-%d %H:%M")
+    if isinstance(date_input, str):
+        date_obj = parser.isoparse(date_input)
+    elif isinstance(date_input, datetime):
+        date_obj = date_input
+    else:
+        date_obj = None
+
+    if date_obj:
+        return date_obj.strftime("%Y-%m-%d %H:%M")
 
     if default_string:
         return default_string

probely/cli/tables/finding_table.py

@@ -1,5 +1,3 @@
-from typing import Dict
-
 from rich.table import Table
 
 from probely.cli.renderers import (
@@ -9,6 +7,7 @@
 )
 from probely.cli.tables.base_table import BaseOutputTable
 from probely.sdk.enums import FindingSeverityEnum
+from probely.sdk._schemas import Finding
 
 DEFAULT_LAST_FOUND_DATE_VALUE = "NO_DATE"
 
@@ -29,19 +28,16 @@ def create_table(cls, show_header: bool = False) -> Table:
         return table
 
     @classmethod
-    def add_row(cls, table: Table, finding: Dict) -> None:
-        target = finding.get("target", {})
-
+    def add_row(cls, table: Table, finding: Finding) -> None:
+        target = finding.target
         table.add_row(
-            f"{target['id']}-{finding['id']}",  # Composite Finding ID
-            target["id"],  # target_id
-            get_printable_enum_value(
-                FindingSeverityEnum, finding["severity"]
-            ),  # severity
-            finding["definition"]["name"],  # title
+            f"{target.id}-{finding.id}",  # Composite Finding ID
+            target.id,
+            get_printable_enum_value(FindingSeverityEnum, finding.severity),  # severity
+            finding.definition.name,  # title
             get_printable_date(
-                finding["last_found"], DEFAULT_LAST_FOUND_DATE_VALUE
+                finding.last_found, DEFAULT_LAST_FOUND_DATE_VALUE
             ),  # last_found
-            finding["state"],  # state
-            get_printable_labels(finding["labels"]),  # labels
+            finding.state,
+            get_printable_labels(finding.labels),
         )