Skip to content

Commit 9ed06ff

Browse files
Vladimir KontićJoão Silva
Vladimir Kontić
authored and
João Silva
committed
!35 Pagination for Scans and project structure improvements Closes #79
1 parent 7f54433 commit 9ed06ff

56 files changed

Lines changed: 630 additions & 637 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎examples/call_targets_sdk.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,16 @@
11
import json
22
from typing import List
33

4-
from probely_cli import list_targets, Probely
4+
from probely_cli import Probely, list_targets
55
from probely_cli.exceptions import ProbelyException
66

77
if __name__ == "__main__":
88

99
# You can config programmatically
1010
Probely.init(api_key="your_api_key")
1111

12-
targets_list: List[dict]
13-
1412
try:
15-
targets_list = list_targets()
13+
targets_list: List[dict] = list_targets()
1614

1715
for target in targets_list:
1816
print(json.dumps(target, indent=4))

‎probely_cli/__init__.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
from .sdk.client import Probely
22
from .sdk.targets import (
3+
add_target,
4+
delete_target,
5+
delete_targets,
36
list_targets,
47
retrieve_target,
58
retrieve_targets,
6-
delete_target,
7-
delete_targets,
8-
add_target,
99
update_target,
1010
update_targets,
1111
)

‎probely_cli/cli/app.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import argparse
22
import logging
33

4-
from probely_cli import settings, Probely
5-
from probely_cli.exceptions import ProbelyException, ProbelyCLIValidation
4+
from probely_cli import Probely, settings
5+
from probely_cli.exceptions import ProbelyCLIValidation, ProbelyException
66

77
logger = logging.getLogger(__name__)
88

‎probely_cli/cli/commands/apply/apply.py‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,7 @@
22

33
from probely_cli.cli.commands.apply.schemas import ApplyFileSchema
44
from probely_cli.cli.common import validate_and_retrieve_yaml_content
5-
from probely_cli.exceptions import (
6-
ProbelyException,
7-
ProbelyBadRequest,
8-
)
5+
from probely_cli.exceptions import ProbelyBadRequest, ProbelyException
96
from probely_cli.sdk.targets import add_target
107

118
logger = logging.getLogger(__name__)

‎probely_cli/cli/commands/apply/schemas.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from marshmallow import Schema, fields, INCLUDE
1+
from marshmallow import INCLUDE, Schema, fields
22

33

44
class ApplyFileSchema(Schema):

‎probely_cli/cli/commands/findings/get.py‎

Lines changed: 17 additions & 112 deletions
Original file line numberDiff line numberDiff line change
@@ -1,123 +1,28 @@
1-
import json
2-
import sys
3-
from typing import Dict, Generator, Optional
4-
5-
import marshmallow
6-
import yaml
7-
from marshmallow import Schema
8-
from rich.console import Console
9-
from rich.live import Live
10-
from rich.table import Table
1+
import argparse
112

123
from probely_cli.cli.commands.findings.schemas import FindingsApiFiltersSchema
13-
from probely_cli.cli.common import OutputEnum
14-
from probely_cli.cli.formatters import (
15-
get_printable_enum_value,
16-
get_printable_date,
17-
get_printable_labels,
18-
)
4+
from probely_cli.cli.common import prepare_filters_for_api
5+
from probely_cli.cli.enums import EntityTypeEnum, OutputEnum
6+
from probely_cli.cli.renderers import OutputRenderer
197
from probely_cli.exceptions import ProbelyCLIValidation
20-
from probely_cli.sdk.common import FindingSeverityEnum
21-
from probely_cli.sdk.findings import list_findings
22-
from probely_cli.sdk.findings import retrieve_findings
23-
24-
25-
DEFAULT_LAST_FOUND_DATE_VALUE = "NO_DATE"
26-
27-
28-
def finding_filters_handler(args):
29-
filters_schema: Schema = FindingsApiFiltersSchema()
30-
try:
31-
api_ready_filters = filters_schema.load(vars(args))
32-
except marshmallow.ValidationError as e:
33-
raise ProbelyCLIValidation(str(e))
34-
35-
return api_ready_filters
36-
37-
38-
def create_findings_table(show_header: bool = False) -> Table:
39-
"""
40-
Initializes and returns a Rich Table for displaying Findings.
41-
"""
42-
table = Table(show_header=show_header, box=None)
43-
44-
table.add_column("ID", width=18)
45-
table.add_column("TARGET_ID", width=12)
46-
table.add_column("SEVERITY", width=8)
47-
table.add_column("TITLE", width=48, no_wrap=True)
48-
table.add_column("LAST_FOUND", width=16)
49-
table.add_column("STATE", width=8)
50-
table.add_column("LABELS", width=16)
51-
52-
return table
53-
54-
55-
def add_finding_to_table(table: Table, finding: Dict) -> None:
56-
"""
57-
Adds a single finding as a row to the provided Rich Table.
58-
"""
59-
target = finding.get("target")
60-
61-
table.add_row(
62-
f"{target['id']}-{finding['id']}", # Composite Finding ID
63-
target["id"], # target_id
64-
get_printable_enum_value(FindingSeverityEnum, finding["severity"]), # severity
65-
finding["definition"]["name"], # title
66-
get_printable_date(
67-
finding["last_found"], DEFAULT_LAST_FOUND_DATE_VALUE
68-
), # last_found
69-
finding["state"], # state
70-
get_printable_labels(finding["labels"]), # labels
71-
)
72-
73-
74-
def render_json_output(findings: Generator[Dict, None, None], console: Console) -> None:
75-
console.print("[")
76-
first = True
77-
for finding in findings:
78-
if not first:
79-
console.print(",")
80-
console.print(json.dumps(finding, indent=2))
81-
first = False
82-
console.print("]")
83-
84-
85-
def render_yaml_output(findings: Generator[Dict, None, None], console: Console) -> None:
86-
for finding in findings:
87-
console.print(yaml.dump([finding], indent=2, width=sys.maxsize))
8+
from probely_cli.sdk.findings import list_findings, retrieve_findings
889

8910

90-
def render_table_output(
91-
findings: Generator[Dict, None, None], console: Console
92-
) -> None:
93-
table = create_findings_table(show_header=True)
94-
with Live(table, refresh_per_second=1, console=console):
95-
for finding in findings:
96-
add_finding_to_table(table, finding)
97-
98-
99-
def render_output(
100-
findings: Generator[Dict, None, None],
101-
output_type: Optional[OutputEnum],
102-
console: Console,
103-
) -> None:
104-
if output_type == OutputEnum.JSON:
105-
render_json_output(findings, console)
106-
elif output_type == OutputEnum.YAML:
107-
render_yaml_output(findings, console)
108-
else:
109-
render_table_output(findings, console)
110-
111-
112-
def findings_get_command_handler(args):
113-
api_filters = finding_filters_handler(args)
114-
if api_filters and args.findings_ids:
115-
raise ProbelyCLIValidation("filters and finding ids are mutually exclusive.")
11+
def findings_get_command_handler(args: argparse.Namespace):
12+
filters = prepare_filters_for_api(FindingsApiFiltersSchema, args)
13+
if filters and args.findings_ids:
14+
raise ProbelyCLIValidation("filters and Finding IDs are mutually exclusive.")
11615

11716
if args.findings_ids:
11817
findings_generator = retrieve_findings(findings_ids=args.findings_ids)
11918
else:
120-
findings_generator = list_findings(findings_filters=api_filters)
19+
findings_generator = list_findings(findings_filters=filters)
12120

12221
output_type = OutputEnum[args.output] if args.output else None
123-
render_output(findings_generator, output_type, args.console)
22+
renderer = OutputRenderer(
23+
records=findings_generator,
24+
output_type=output_type,
25+
console=args.console,
26+
entity_type=EntityTypeEnum.FINDING,
27+
)
28+
renderer.render()

‎probely_cli/cli/commands/findings/parsers.py‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,9 @@
33
from rich_argparse import RichHelpFormatter
44

55
from probely_cli.cli.commands.findings.get import findings_get_command_handler
6-
from probely_cli.cli.common import (
7-
show_help,
8-
)
9-
from probely_cli.sdk.common import FindingSeverityEnum, FindingStateEnum
10-
from probely_cli.settings import TRUTHY_VALUES, FALSY_VALUES
6+
from probely_cli.cli.common import show_help
7+
from probely_cli.sdk.enums import FindingSeverityEnum, FindingStateEnum
8+
from probely_cli.settings import FALSY_VALUES, TRUTHY_VALUES
119

1210

1311
def build_findings_filters_parser() -> argparse.ArgumentParser:

‎probely_cli/cli/commands/findings/schemas.py‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,8 @@
11
import marshmallow
22

3-
from probely_cli.cli.common import (
4-
ProbelyCLIEnumField,
5-
ProbelyCLIBaseFiltersSchema,
6-
)
7-
from probely_cli.sdk.common import FindingSeverityEnum, FindingStateEnum
3+
from probely_cli.cli.schemas.base_schema import ProbelyCLIBaseFiltersSchema
4+
from probely_cli.cli.schemas.custom_fields import ProbelyCLIEnumField
5+
from probely_cli.sdk.enums import FindingSeverityEnum, FindingStateEnum
86

97

108
class FindingsApiFiltersSchema(ProbelyCLIBaseFiltersSchema):

‎probely_cli/cli/commands/parsers.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,8 @@
66
from probely_cli.cli.commands.findings.parsers import build_findings_parser
77
from probely_cli.cli.commands.scans.parsers import build_scans_parser
88
from probely_cli.cli.commands.targets.parsers import build_targets_parser
9-
from probely_cli.cli.common import (
10-
OutputEnum,
11-
show_help,
12-
)
9+
from probely_cli.cli.common import show_help
10+
from probely_cli.cli.enums import OutputEnum
1311
from probely_cli.version import __version__
1412

1513

‎probely_cli/cli/commands/scans/cancel.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,19 @@
11
import logging
2-
from probely_cli.cli.commands.scans.get import (
3-
prepare_scan_filters_for_api,
2+
3+
from probely_cli.cli.commands.scans.schemas import ScanApiFiltersSchema
4+
from probely_cli.cli.common import (
5+
display_scans_response_output,
6+
prepare_filters_for_api,
47
)
5-
from probely_cli.cli.common import display_scans_response_output
68
from probely_cli.exceptions import ProbelyCLIValidation
7-
from probely_cli.sdk.scans import cancel_scans, cancel_scan, list_scans
9+
from probely_cli.sdk.scans import cancel_scan, cancel_scans, list_scans
810

911
logger = logging.getLogger(__name__)
1012

1113

1214
def scans_cancel_command_handler(args):
1315
scan_ids = args.scan_ids
14-
filters = prepare_scan_filters_for_api(args)
16+
filters = prepare_filters_for_api(ScanApiFiltersSchema, args)
1517

1618
if not scan_ids and not filters:
1719
raise ProbelyCLIValidation("Expected scan_ids or filters")

0 commit comments

Comments
 (0)