Skip to content

Commit 4dd9fe1

Browse files
author
Juliya Smith
authored
Bugfix/json output (#143)
1 parent 2ecbd62 commit 4dd9fe1

13 files changed

Lines changed: 579 additions & 361 deletions

‎src/code42cli/cmds/alert_rules.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
from code42cli.options import format_option
1616
from code42cli.options import OrderedGroup
1717
from code42cli.options import sdk_options
18-
from code42cli.output_formats import get_output_format_func
18+
from code42cli.output_formats import OutputFormatter
1919

2020

2121
class AlertRuleTypes:
@@ -78,11 +78,11 @@ def remove_user(state, rule_id, username):
7878
@sdk_options()
7979
def list_alert_rules(state, format=None):
8080
"""Fetch existing alert rules."""
81-
format_func = get_output_format_func(format)
81+
formatter = OutputFormatter(format, _HEADER_KEYS_MAP)
8282
selected_rules = _get_all_rules_metadata(state.sdk)
8383
if selected_rules:
84-
formatted_output = format_func(selected_rules, _HEADER_KEYS_MAP)
85-
echo(formatted_output)
84+
for output in formatter.get_formatted_output(selected_rules):
85+
echo(output)
8686

8787

8888
@alert_rules.command()

‎src/code42cli/cmds/alerts.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import code42cli.options as opt
1313
from code42cli.cmds.search.cursor_store import AlertCursorStore
1414
from code42cli.options import format_option
15-
from code42cli.output_formats import get_output_format_func
15+
from code42cli.output_formats import OutputFormatter
1616

1717
SEARCH_DEFAULT_HEADER = OrderedDict()
1818
SEARCH_DEFAULT_HEADER["name"] = "RuleName"
@@ -185,15 +185,15 @@ def search(
185185
output_header = ext.try_get_default_header(
186186
include_all, SEARCH_DEFAULT_HEADER, format
187187
)
188-
format_func = get_output_format_func(format)
188+
formatter = OutputFormatter(format, output_header)
189189
cursor = _get_alert_cursor_store(cli_state.profile.name) if use_checkpoint else None
190190
handlers = ext.create_handlers(
191191
cli_state.sdk,
192192
AlertExtractor,
193193
cursor,
194194
use_checkpoint,
195-
format_func=format_func,
196-
output_header=output_header,
195+
formatter=formatter,
196+
force_pager=include_all,
197197
)
198198
extractor = _get_alert_extractor(cli_state.sdk, handlers)
199199
extractor.use_or_query = or_query

‎src/code42cli/cmds/legal_hold.py‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
from code42cli.options import format_option
1919
from code42cli.options import OrderedGroup
2020
from code42cli.options import sdk_options
21-
from code42cli.output_formats import get_output_format_func
21+
from code42cli.output_formats import OutputFormatter
2222
from code42cli.util import format_string_list_to_columns
2323

2424

@@ -76,11 +76,11 @@ def remove_user(state, matter_id, username):
7676
@sdk_options()
7777
def _list(state, format=None):
7878
"""Fetch existing legal hold matters."""
79-
format_func = get_output_format_func(format)
79+
formatter = OutputFormatter(format, _MATTER_KEYS_MAP)
8080
matters = _get_all_active_matters(state.sdk)
8181
if matters:
82-
output = format_func(matters, _MATTER_KEYS_MAP)
83-
echo(output)
82+
for output in formatter.get_formatted_output(matters):
83+
echo(output)
8484

8585

8686
@legal_hold.command()
@@ -100,9 +100,10 @@ def _list(state, format=None):
100100
@sdk_options()
101101
def show(state, matter_id, include_inactive=False, include_policy=False, format=None):
102102
"""Display details of a given legal hold matter."""
103-
format_func = get_output_format_func(format)
103+
formatter = OutputFormatter(format, _MATTER_KEYS_MAP)
104104
matter = _check_matter_is_accessible(state.sdk, matter_id)
105105
matter["creator_username"] = matter["creator"]["username"]
106+
matter = json.loads(matter.text)
106107

107108
# if `active` is None then all matters (whether active or inactive) are returned. True returns
108109
# only those that are active.
@@ -117,8 +118,8 @@ def show(state, matter_id, include_inactive=False, include_policy=False, format=
117118
member["user"]["username"] for member in memberships if not member["active"]
118119
]
119120

120-
output = format_func([matter], _MATTER_KEYS_MAP)
121-
echo(output)
121+
for output in formatter.get_formatted_output([matter]):
122+
echo(output)
122123

123124
_print_matter_members(active_usernames, member_type="active")
124125

‎src/code42cli/cmds/search/extraction.py‎

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,9 @@
1919
def try_get_default_header(include_all, default_header, output_format):
2020
"""Returns appropriate header based on include-all and output format. If returns None,
2121
the CLI format option will figure out the header based on the data keys."""
22-
output_header = None
23-
24-
if output_format == OutputFormat.TABLE and not include_all:
25-
output_header = default_header
26-
elif output_format != OutputFormat.TABLE and include_all:
27-
err_text = "--include-all only allowed for non-Table output formats."
22+
output_header = None if include_all else default_header
23+
if output_format != OutputFormat.TABLE and include_all:
24+
err_text = "--include-all only allowed for Table output format."
2825
logger.log_error(err_text)
2926
raise errors.Code42CLIError(err_text)
3027
return output_header
@@ -45,7 +42,7 @@ def _get_alert_details(sdk, alert_summary_list):
4542

4643

4744
def create_handlers(
48-
sdk, extractor_class, cursor_store, checkpoint_name, format_func, output_header,
45+
sdk, extractor_class, cursor_store, checkpoint_name, formatter, force_pager,
4946
):
5047
extractor = extractor_class(sdk, ExtractionHandlers())
5148
handlers = ExtractionHandlers()
@@ -83,14 +80,16 @@ def handle_response(response):
8380
total_events = len(events)
8481
handlers.TOTAL_EVENTS += total_events
8582

86-
def paginate():
87-
yield format_func(events, output_header)
83+
def _format_output():
84+
return formatter.get_formatted_output(events)
8885

89-
if len(events) > 10:
90-
click.echo_via_pager(paginate)
86+
if len(events) > 10 or force_pager:
87+
click.echo_via_pager(_format_output())
9188
else:
92-
for page in paginate():
93-
click.echo(page)
89+
for page in _format_output():
90+
click.echo(page, nl=False)
91+
if formatter.output_format == OutputFormat.TABLE:
92+
click.echo()
9493

9594
# To make sure the extractor records correct timestamp event when `CTRL-C` is pressed.
9695
if total_events:

‎src/code42cli/cmds/securitydata.py‎

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,13 @@
1111
import code42cli.cmds.search.options as searchopt
1212
import code42cli.errors as errors
1313
from code42cli.cmds.search.cursor_store import FileEventCursorStore
14-
from code42cli.cmds.securitydata_output_formats import (
15-
get_file_events_output_format_func,
16-
)
14+
from code42cli.cmds.securitydata_output_formats import FileEventsOutputFormatter
1715
from code42cli.logger import get_main_cli_logger
1816
from code42cli.options import format_option
1917
from code42cli.options import incompatible_with
2018
from code42cli.options import OrderedGroup
2119
from code42cli.options import sdk_options
22-
from code42cli.output_formats import get_output_format_func
20+
from code42cli.output_formats import OutputFormatter
2321

2422
logger = get_main_cli_logger()
2523

@@ -210,7 +208,7 @@ def search(
210208
output_header = ext.try_get_default_header(
211209
include_all, SEARCH_DEFAULT_HEADER, format
212210
)
213-
format_func = get_file_events_output_format_func(format)
211+
formatter = FileEventsOutputFormatter(format, output_header)
214212
cursor = (
215213
_get_file_event_cursor_store(state.profile.name) if use_checkpoint else None
216214
)
@@ -219,8 +217,8 @@ def search(
219217
FileEventExtractor,
220218
cursor,
221219
use_checkpoint,
222-
format_func=format_func,
223-
output_header=output_header,
220+
formatter=formatter,
221+
force_pager=include_all,
224222
)
225223
extractor = _get_file_event_extractor(state.sdk, handlers)
226224
extractor.use_or_query = or_query
@@ -250,12 +248,12 @@ def saved_search(state):
250248
@sdk_options()
251249
def _list(state, format=None):
252250
"""List available saved searches."""
253-
format_func = get_output_format_func(format)
251+
formatter = OutputFormatter(format, _HEADER_KEYS_MAP)
254252
response = state.sdk.securitydata.savedsearches.get()
255-
result = response["searches"]
256-
if result:
257-
output = format_func(result, _HEADER_KEYS_MAP)
258-
echo(output)
253+
saved_searches = response["searches"]
254+
if saved_searches:
255+
for output in formatter.get_formatted_output(saved_searches):
256+
echo(output)
259257

260258

261259
@saved_search.command()

‎src/code42cli/cmds/securitydata_output_formats.py‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,17 +9,24 @@
99
import code42cli.cmds.search.enums as enum
1010
from code42cli.output_formats import CEF_DEFAULT_PRODUCT_NAME
1111
from code42cli.output_formats import CEF_DEFAULT_SEVERITY_LEVEL
12-
from code42cli.output_formats import get_output_format_func
12+
from code42cli.output_formats import OutputFormatter
1313

1414

15-
def get_file_events_output_format_func(value):
16-
if value == enum.FileEventsOutputFormat.CEF:
17-
return to_cef
18-
return get_output_format_func(value)
15+
class FileEventsOutputFormatter(OutputFormatter):
16+
def __init__(self, output_format, header=None):
17+
output_format = (
18+
output_format.upper()
19+
if output_format
20+
else enum.FileEventsOutputFormat.TABLE
21+
)
22+
super().__init__(output_format, header)
23+
if output_format == enum.FileEventsOutputFormat.CEF:
24+
self._format_func = to_cef
1925

2026

21-
def to_cef(output, header):
22-
return [_convert_event_to_cef(e) for e in output]
27+
def to_cef(output):
28+
"""Output is a single record"""
29+
return "{}\n".format(_convert_event_to_cef(output))
2330

2431

2532
def _convert_event_to_cef(event):

‎src/code42cli/output_formats.py‎

Lines changed: 46 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -20,54 +20,66 @@ def __iter__(self):
2020
return iter([self.TABLE, self.CSV, self.JSON, self.RAW])
2121

2222

23-
def get_output_format_func(value):
24-
if value is not None:
25-
value = value.upper()
26-
if value == OutputFormat.CSV:
27-
return to_csv
28-
if value == OutputFormat.RAW:
29-
return to_json
30-
if value == OutputFormat.TABLE:
31-
return to_table
32-
if value == OutputFormat.JSON:
33-
return to_formatted_json
34-
# default option
35-
return to_table
36-
37-
38-
def to_csv(output, header):
23+
class OutputFormatter:
24+
def __init__(self, output_format, header=None):
25+
output_format = output_format.upper() if output_format else OutputFormat.TABLE
26+
self.output_format = output_format
27+
self._format_func = to_table
28+
self.header = header
29+
30+
if output_format == OutputFormat.CSV:
31+
self._format_func = to_csv
32+
elif output_format == OutputFormat.RAW:
33+
self._format_func = to_json
34+
elif output_format == OutputFormat.TABLE:
35+
self._format_func = self._to_table
36+
elif output_format == OutputFormat.JSON:
37+
self._format_func = to_formatted_json
38+
39+
def _format_output(self, output):
40+
return self._format_func(output)
41+
42+
def _to_table(self, output):
43+
return to_table(output, self.header)
44+
45+
def get_formatted_output(self, output):
46+
if self._requires_list_output:
47+
yield self._format_output(output)
48+
else:
49+
for item in output:
50+
yield self._format_output(item)
51+
52+
@property
53+
def _requires_list_output(self):
54+
return self.output_format in (OutputFormat.TABLE, OutputFormat.CSV)
55+
56+
57+
def to_csv(output):
58+
"""Output is a list of records"""
3959
if not output:
4060
return
4161
string_io = io.StringIO()
42-
writer = csv.DictWriter(string_io, fieldnames=output[0].keys())
62+
fieldnames = list({k for d in output for k in d.keys()})
63+
writer = csv.DictWriter(string_io, fieldnames=fieldnames)
4364
writer.writeheader()
4465
writer.writerows(output)
4566
return string_io.getvalue()
4667

4768

4869
def to_table(output, header):
70+
"""Output is a list of records"""
4971
if not output:
5072
return
51-
header = header or get_dynamic_header(output[0])
5273
rows, column_size = find_format_width(output, header)
5374
return format_to_table(rows, column_size)
5475

5576

56-
def _filter(output, header):
57-
return [{header[key]: row[key] for key in header.keys()} for row in output]
77+
def to_json(output):
78+
"""Output is a single record"""
79+
return "{}\n".format(json.dumps(output))
5880

5981

60-
def to_json(output, header=None):
61-
return json.dumps(output)
62-
63-
64-
def to_formatted_json(output, header=None):
65-
return json.dumps(_filter(output, header), indent=4)
66-
67-
68-
def get_dynamic_header(header_items):
69-
return {
70-
key: key.capitalize()
71-
for key in header_items.keys()
72-
if type(header_items[key]) == str
73-
}
82+
def to_formatted_json(output):
83+
"""Output is a single record"""
84+
json_str = "{}\n".format(json.dumps(output, indent=4))
85+
return json_str

‎src/code42cli/util.py‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -44,20 +44,21 @@ def find_format_width(record, header, include_header=True):
4444
Finds the largest string against each column so as to decide the padding size for the column.
4545
4646
Args:
47-
record (list of dict), data to be formatted.
48-
header (dict), key-value where keys should map to keys of record dict and
47+
record (dict): data to be formatted.
48+
header (dict): key-value where keys should map to keys of record dict and
4949
value is the corresponding column name to be displayed on the CLI.
50-
include_header (bool), include header in output, defaults to True.
50+
include_header (bool): include header in output, defaults to True.
5151
5252
Returns:
53-
tuple (list of dict, dict), i.e Filtered records, padding size of columns.
53+
tuple (list of dict, dict): i.e Filtered records, padding size of columns.
5454
"""
5555
rows = []
5656
if include_header:
57+
if not header:
58+
header = _get_default_header(record)
5759
rows.append(header)
5860

59-
# Set default max width items to column names
60-
max_width_item = dict(header.items())
61+
max_width_item = dict(header.items()) # Copy
6162
for record_row in record:
6263
row = OrderedDict()
6364
for header_key in header.keys():
@@ -144,3 +145,17 @@ def inner(*args, **kwargs):
144145
return func(*args, **kwargs)
145146

146147
return inner
148+
149+
150+
def _get_default_header(header_items):
151+
if not header_items:
152+
return
153+
154+
# Creates dict where keys and values are the same for `find_format_width()`.
155+
header = {}
156+
for item in header_items:
157+
keys = item.keys()
158+
for key in keys:
159+
if key not in header and isinstance(key, str):
160+
header[key] = key
161+
return header

0 commit comments

Comments
 (0)