From 801c233ec38eb9d25e02358320ac4efebfa3028f Mon Sep 17 00:00:00 2001 From: Peter Kwan Date: Wed, 8 Jan 2025 17:14:59 -0800 Subject: [PATCH] fixed a bug where the offset isn't displayed for query yielding more results than the first page; also added support for the global and desc flag for search-by-term --- praetorian_cli/handlers/search.py | 9 ++++++--- praetorian_cli/handlers/utils.py | 2 +- praetorian_cli/sdk/chariot.py | 5 +++-- praetorian_cli/sdk/entities/search.py | 6 +++++- praetorian_cli/sdk/model/globals.py | 17 +++++++++-------- praetorian_cli/sdk/test/test_z_cli.py | 3 +++ 6 files changed, 27 insertions(+), 15 deletions(-) diff --git a/praetorian_cli/handlers/search.py b/praetorian_cli/handlers/search.py index 268b5b0..f015b12 100644 --- a/praetorian_cli/handlers/search.py +++ b/praetorian_cli/handlers/search.py @@ -11,7 +11,9 @@ @click.option('-t', '--term', help='Enter a search term', required=True) @click.option('-c', '--count', is_flag=True, default=False, help='Return statistics on search') @click.option('-d', '--details', is_flag=True, default=False, help='Show detailed information') -def search(chariot, term, count, details, offset, page): +@click.option('-desc', '--desc', is_flag=True, default=False, help='Return data in descending order') +@click.option('-g', '--global', 'global_', is_flag=True, default=False, help='Use the global data set') +def search(chariot, term, count, details, offset, page, desc, global_): """ Query Chariot for matches or counts using the search syntax \b @@ -53,9 +55,10 @@ def search(chariot, term, count, details, offset, page): - praetorian chariot search --term "status:OH" - praetorian chariot search --term "status:OH" --details --page all - praetorian chariot search --term "#asset#www.example.com" - - praetorian chariot search --term "dns:https://github.com/praetorian-inc/" + - praetorian chariot search --term "dns:https://github.com/praetorian-inc/" --desc """ if count: print_json(chariot.search.count(term)) else: - render_list_results(chariot.search.by_term(term, offset, pagination_size(page)), details) + render_list_results(chariot.search.by_term(term, offset, pagination_size(page), desc=desc, global_=global_), + details) diff --git a/praetorian_cli/handlers/utils.py b/praetorian_cli/handlers/utils.py index 3b68dab..71ffc01 100644 --- a/praetorian_cli/handlers/utils.py +++ b/praetorian_cli/handlers/utils.py @@ -24,7 +24,7 @@ def render_list_results(list_results, details): def render_offset(offset): if offset: click.echo('There are more results. Add the following argument to the command to view them:') - click.echo(f'--offset "{json.dumps(offset)}"') + click.echo(f'--offset {json.dumps(offset)}') def pagination_size(page): diff --git a/praetorian_cli/sdk/chariot.py b/praetorian_cli/sdk/chariot.py index d3a6e1a..6c1fb41 100644 --- a/praetorian_cli/sdk/chariot.py +++ b/praetorian_cli/sdk/chariot.py @@ -41,12 +41,13 @@ def my(self, params: dict, pages=1) -> {}: process_failure(resp) resp = resp.json() extend(final_resp, resp) - if 'offset' not in resp: break - params['offset'] = json.dumps(resp['offset']) + if 'offset' in resp: + final_resp['offset'] = json.dumps(resp['offset']) + return final_resp def post(self, type: str, params): diff --git a/praetorian_cli/sdk/entities/search.py b/praetorian_cli/sdk/entities/search.py index 256f692..1d50965 100644 --- a/praetorian_cli/sdk/entities/search.py +++ b/praetorian_cli/sdk/entities/search.py @@ -32,12 +32,16 @@ def by_ip(self, ip_prefix, offset=None, pages=10000) -> tuple: def by_dns(self, dns_prefix, offset=None, pages=10000) -> tuple: return self.by_term(f'dns:{dns_prefix}', offset, pages) - def by_term(self, search_term, offset=None, pages=1000, exact=False) -> tuple: + def by_term(self, search_term, offset=None, pages=1000, exact=False, desc=False, global_=False) -> tuple: params = dict(key=search_term) if offset: params = params | dict(offset=offset) if exact: params = params | dict(exact='true') + if desc: + params = params | dict(desc='true') + if global_: + params = params | {'global': 'true'} # extract all the different types of entities in the search results into a # flattened list of `hits` diff --git a/praetorian_cli/sdk/model/globals.py b/praetorian_cli/sdk/model/globals.py index f8ca357..17c5994 100644 --- a/praetorian_cli/sdk/model/globals.py +++ b/praetorian_cli/sdk/model/globals.py @@ -83,29 +83,32 @@ class AddRisk(Enum): CAPABILITIES = ( + 'reverse-whois', + 'csp-mine', + 'tls-mine', + 'azuread-discovery', + 'edgar', + 'cidr', 'favicon', 'reverse-csp', + 'builtwith', 'nuclei', 'whois', 'subdomain', - 'csp-mine', - 'tls-mine', 'portscan', 'github', 'github-repository', 'secrets', 'amazon', + 'bitbucket', 'azure', 'gcp', 'ns1', 'cloudflare', 'gato', - 'crowdstrike', 'crawler', 'gitlab', 'ssh', - 'azuread-discovery', - 'edgar', 'nessus', 'nessus-import', 'insightvm', @@ -114,11 +117,9 @@ class AddRisk(Enum): 'qualys-import', 'burp-enterprise', 'ip', - 'cidr', 'website', - 'reverse-whois', 'digitalocean', 'burp-internal', 'seed-import', - 'builtwith' + 'tenablevm', ) diff --git a/praetorian_cli/sdk/test/test_z_cli.py b/praetorian_cli/sdk/test/test_z_cli.py index 4fcbb9d..742c6a3 100644 --- a/praetorian_cli/sdk/test/test_z_cli.py +++ b/praetorian_cli/sdk/test/test_z_cli.py @@ -139,6 +139,9 @@ def test_search_cli(self): self.verify(f'add asset -n {o.asset_name} -d {o.asset_dns}') self.verify(f'search -t "#asset#{o.asset_dns}" -p all', [o.asset_key]) + self.verify(f'search -t "#asset#{o.asset_dns}" -p all --desc', [o.asset_key]) + self.verify(f'search -t "#asset#{o.asset_dns}" -p all -g') + self.verify(f'search -t "#asset#{o.asset_dns}" -d -p all', [o.asset_key, '"key"', '"data"']) self.verify(f'search -t "#asset#{o.asset_dns}" -c -p all', ['"A": 1'])