!40 78 tagets pagination and generators Closes #78 and #49

Author: Vladimir Kontić

examples/call_targets_sdk.py

@@ -1,17 +1,19 @@
 import json
-from typing import List
 
 from probely_cli import Probely, ProbelyException, list_targets
+from probely_cli.exceptions import ProbelyRequestFailed
 
 if __name__ == "__main__":
 
     Probely.init(api_key="your_api_key")
 
     try:
-        targets_list: List[dict] = list_targets()
-
-        for target in targets_list:
+        targets_generator = list_targets()
+        for target in targets_generator():
             print(json.dumps(target, indent=4))
-
+    except ProbelyRequestFailed as request_error:
+        print("Request to Probely API failed:", request_error)
     except ProbelyException as probely_exception:
-        print("Probely sdk error:", probely_exception)
+        print("Probely SDK error:", probely_exception)
+    except Exception as general_error:
+        print("An unexpected error occurred:", general_error)

probely_cli/cli/commands/targets/delete.py

@@ -1,4 +1,5 @@
-from probely_cli.cli.commands.targets.get import target_filters_handler
+from probely_cli.cli.commands.targets.schemas import TargetApiFiltersSchema
+from probely_cli.cli.common import prepare_filters_for_api
 from probely_cli.exceptions import ProbelyCLIValidation
 from probely_cli.sdk.targets import delete_target, delete_targets, list_targets
 
@@ -7,7 +8,7 @@ def targets_delete_command_handler(args):
     """
     Delete targets
     """
-    filters = target_filters_handler(args)
+    filters = prepare_filters_for_api(TargetApiFiltersSchema, args)
     targets_ids = args.target_ids
 
     if not filters and not targets_ids:
@@ -17,11 +18,15 @@ def targets_delete_command_handler(args):
         raise ProbelyCLIValidation("filters and Target IDs are mutually exclusive.")
 
     if filters:
-        searched_targets = list_targets(targets_filters=filters)
-        if not searched_targets:
+        generator = list_targets(targets_filters=filters)
+        first_target = next(generator, None)
+
+        if not first_target:
             raise ProbelyCLIValidation("Selected Filters returned no results")
 
-        targets_ids = [target.get("id") for target in searched_targets]
+        targets_ids = [first_target.get("id")] + [
+            target.get("id") for target in generator
+        ]
 
     if len(targets_ids) == 1:
         target_id = delete_target(targets_ids[0])

probely_cli/cli/commands/targets/get.py

@@ -1,102 +1,33 @@
-import json
-import sys
-from typing import Dict, List, Union
-
-import marshmallow
-import yaml
-from marshmallow import Schema
-from rich.table import Table
-
 from probely_cli.cli.commands.targets.schemas import TargetApiFiltersSchema
-from probely_cli.cli.enums import OutputEnum
-from probely_cli.cli.renderers import (
-    get_printable_date,
-    get_printable_enum_value,
-    get_printable_labels,
-)
+from probely_cli.cli.common import prepare_filters_for_api
+from probely_cli.cli.enums import EntityTypeEnum, OutputEnum
+
+from probely_cli.cli.renderers import OutputRenderer
 from probely_cli.exceptions import ProbelyCLIValidation
-from probely_cli.sdk.enums import TargetRiskEnum
 from probely_cli.sdk.targets import list_targets, retrieve_targets
 
-TARGET_NEVER_SCANNED_OUTPUT: str = "Never_scanned"
-
-
-def _get_printable_last_scan_date(target: Dict) -> str:
-    last_scan_obj: Union[dict, None] = target.get("last_scan", None)
-
-    if last_scan_obj is None:
-        return TARGET_NEVER_SCANNED_OUTPUT
-
-    last_scan_start_date_str: Union[str, None] = last_scan_obj.get("started", None)
-
-    return get_printable_date(last_scan_start_date_str, TARGET_NEVER_SCANNED_OUTPUT)
-
-
-def get_tabled_targets(targets_list: List[Dict]) -> Table:
-    table = Table(box=None)
-    table.add_column("ID")
-    table.add_column("NAME")
-    table.add_column("URL")
-    table.add_column("RISK")
-    table.add_column("LAST_SCAN")
-    table.add_column("LABELS")
-
-    for target in targets_list:
-        asset = target.get("site")
-
-        table.add_row(
-            target.get("id"),
-            asset.get("name", "N/D"),
-            asset.get("url"),
-            get_printable_enum_value(TargetRiskEnum, target["risk"]),
-            _get_printable_last_scan_date(target),  # last_scan
-            get_printable_labels(target["labels"]),
-        )
-
-    return table
-
-
-def target_filters_handler(args) -> dict:
-    filters_schema: Schema = TargetApiFiltersSchema()
-    try:
-        api_ready_filters = filters_schema.load(vars(args))
-    except marshmallow.ValidationError as e:
-        # TODO: translate validations?
-        raise ProbelyCLIValidation(str(e))
-
-    return api_ready_filters
-
-
-def build_cmd_output(args, targets_list):
-    output_type = OutputEnum[args.output] if args.output else None
-
-    if output_type == OutputEnum.JSON:
-        return json.dumps(targets_list, indent=2)
-
-    if output_type == OutputEnum.YAML:
-        return yaml.dump(
-            targets_list,
-            indent=2,
-            width=sys.maxsize,  # avoids word wrapping
-        )
-
-    return get_tabled_targets(targets_list)
-
 
 def targets_get_command_handler(args):
     """
     Lists all accessible targets of client
     """
-    filters = target_filters_handler(args)
+    filters = prepare_filters_for_api(TargetApiFiltersSchema, args)
     targets_ids = args.target_ids
 
     if filters and targets_ids:
         raise ProbelyCLIValidation("filters and Target IDs are mutually exclusive.")
 
     if targets_ids:
-        targets_list = retrieve_targets(targets_ids=targets_ids)
+        targets_generator = retrieve_targets(targets_ids=targets_ids)
     else:
-        targets_list = list_targets(targets_filters=filters)
+        targets_generator = list_targets(targets_filters=filters)
+
+    output_type = OutputEnum[args.output] if args.output else None
 
-    cmd_output = build_cmd_output(args, targets_list)
-    args.console.print(cmd_output)
+    renderer = OutputRenderer(
+        records=targets_generator,
+        output_type=output_type,
+        console=args.console,
+        entity_type=EntityTypeEnum.TARGET,
+    )
+    renderer.render()

probely_cli/cli/commands/targets/start_scan.py

@@ -1,8 +1,9 @@
 import logging
 
-from probely_cli.cli.commands.targets.get import target_filters_handler
+from probely_cli.cli.commands.targets.schemas import TargetApiFiltersSchema
 from probely_cli.cli.common import (
     display_scans_response_output,
+    prepare_filters_for_api,
     validate_and_retrieve_yaml_content,
 )
 from probely_cli.exceptions import ProbelyCLIValidation
@@ -25,7 +26,7 @@ def validate_and_retrieve_extra_payload(args):
 
 
 def start_scans_command_handler(args):
-    filters = target_filters_handler(args)
+    filters = prepare_filters_for_api(TargetApiFiltersSchema, args)
     targets_ids = args.target_ids
 
     if not filters and not targets_ids:
@@ -37,12 +38,13 @@ def start_scans_command_handler(args):
     extra_payload = validate_and_retrieve_extra_payload(args)
 
     if filters:
-        searched_targets = list_targets(targets_filters=filters)
+        generator = list_targets(targets_filters=filters)
+        first_target = next(generator, None)
 
-        if not searched_targets:
+        if not first_target:
             raise ProbelyCLIValidation("Selected Filters returned no results")
 
-        targets_ids = [target["id"] for target in searched_targets]
+        targets_ids = [first_target["id"]] + [target["id"] for target in generator]
 
     if len(targets_ids) == 1:
         scans = [start_scan(targets_ids[0], extra_payload)]

probely_cli/cli/commands/targets/update.py

@@ -5,8 +5,11 @@
 
 import yaml
 
-from probely_cli.cli.commands.targets.get import target_filters_handler
-from probely_cli.cli.common import validate_and_retrieve_yaml_content
+from probely_cli.cli.commands.targets.schemas import TargetApiFiltersSchema
+from probely_cli.cli.common import (
+    prepare_filters_for_api,
+    validate_and_retrieve_yaml_content,
+)
 from probely_cli.cli.enums import OutputEnum
 from probely_cli.exceptions import ProbelyCLIValidation
 from probely_cli.sdk.targets import list_targets, update_target, update_targets
@@ -45,7 +48,7 @@ def update_targets_command_handler(args):
         )
     payload = validate_and_retrieve_yaml_content(yaml_file_path)
 
-    filters = target_filters_handler(args)
+    filters = prepare_filters_for_api(TargetApiFiltersSchema, args)
     targets_ids = args.target_ids
 
     if not filters and not targets_ids:
@@ -65,16 +68,16 @@ def update_targets_command_handler(args):
         return
 
     # Fetch all Targets that match the filters and update them
-    targets_for_update = list_targets(targets_filters=filters)
-    searched_targets_ids = [target["id"] for target in targets_for_update]
+    generator = list_targets(targets_filters=filters)
+    first_target = next(generator, None)
 
-    if not searched_targets_ids:  # TODO test
+    if not first_target:
         raise ProbelyCLIValidation("Selected Filters returned no results")
 
-    target_ids = searched_targets_ids
+    targets_ids = [first_target["id"]] + [target["id"] for target in generator]
 
-    if len(target_ids) == 1:
-        updated_targets = [update_target(target_ids[0], payload)]
+    if len(targets_ids) == 1:
+        updated_targets = [update_target(targets_ids[0], payload)]
     else:
-        updated_targets = update_targets(target_ids, payload)
+        updated_targets = update_targets(targets_ids, payload)
     display_cmd_output(args, updated_targets)

probely_cli/cli/renderers.py

@@ -1,7 +1,14 @@
 import json
 import sys
 import textwrap
-from typing import Dict, Generator, List, Optional, Type, Union
+from typing import (
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Type,
+    Union,
+)
 
 import yaml
 from dateutil import parser
@@ -12,6 +19,7 @@
 from probely_cli.sdk.enums import ProbelyCLIEnum
 
 UNKNOWN_VALUE_REP = "UNKNOWN"
+TARGET_NEVER_SCANNED_OUTPUT = "Never_scanned"
 
 
 class OutputRenderer:
@@ -104,3 +112,14 @@ def get_printable_date(
         return default_string
 
     return ""
+
+
+def get_printable_last_scan_date(target: Dict) -> str:
+    last_scan_obj: Union[dict, None] = target.get("last_scan", None)
+
+    if last_scan_obj is None:
+        return TARGET_NEVER_SCANNED_OUTPUT
+
+    last_scan_start_date_str: Union[str, None] = last_scan_obj.get("started", None)
+
+    return get_printable_date(last_scan_start_date_str, TARGET_NEVER_SCANNED_OUTPUT)

probely_cli/cli/tables/finding_table.py

@@ -1,3 +1,4 @@
+from typing import Dict
 from rich.table import Table
 
 from probely_cli.cli.renderers import (
@@ -25,7 +26,7 @@ def create_table(self, show_header: bool = False) -> Table:
 
         return table
 
-    def add_row(self, table: Table, finding: dict) -> None:
+    def add_row(self, table: Table, finding: Dict) -> None:
         target = finding.get("target", {})
 
         table.add_row(

probely_cli/cli/tables/table_factory.py

@@ -4,6 +4,7 @@
 from probely_cli.cli.tables.base_table import BaseOutputTable
 from probely_cli.cli.tables.finding_table import FindingTable
 from probely_cli.cli.tables.scan_table import ScanTable
+from probely_cli.cli.tables.targets_table import TargetTable
 
 
 class TableFactory:
@@ -12,6 +13,7 @@ def get_table_class(entity_type: EntityTypeEnum) -> Type[BaseOutputTable]:
         ENTITY_TABLE_MAPPING = {
             EntityTypeEnum.FINDING: FindingTable,
             EntityTypeEnum.SCAN: ScanTable,
+            EntityTypeEnum.TARGET: TargetTable,
         }
 
         return ENTITY_TABLE_MAPPING[entity_type]

probely_cli/cli/tables/targets_table.py

@@ -0,0 +1,36 @@
+from typing import Dict
+
+from rich.table import Table
+from probely_cli.cli.renderers import (
+    get_printable_enum_value,
+    get_printable_labels,
+    get_printable_last_scan_date,
+)
+from probely_cli.cli.tables.base_table import BaseOutputTable
+from probely_cli.sdk.enums import TargetRiskEnum
+
+
+class TargetTable(BaseOutputTable):
+    def create_table(self, show_header: bool = False) -> Table:
+        table = Table(show_header=show_header, box=None)
+
+        table.add_column("ID", width=12)
+        table.add_column("NAME", width=36, no_wrap=True)
+        table.add_column("URL", width=48, no_wrap=True)
+        table.add_column("RISK", width=7)
+        table.add_column("LAST_SCAN", width=16)
+        table.add_column("LABELS", width=16, no_wrap=True)
+
+        return table
+
+    def add_row(self, table: Table, target: Dict) -> None:
+        site = target.get("site")
+
+        table.add_row(
+            target.get("id"),
+            site.get("name", "N/D"),
+            site.get("url"),
+            get_printable_enum_value(TargetRiskEnum, target["risk"]),
+            get_printable_last_scan_date(target),
+            get_printable_labels(target["labels"]),
+        )