surveyor/common.py at master · theonlykernel/surveyor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
from pprint import pprint

import click
from cbapi import CbEnterpriseResponseAPI, CbThreatHunterAPI

from products import vmware_cb_response as cbr, vmware_cb_enterprise_edr as cbth


class EDRCommon:
    def __init__(self, product, profile):
        self.product = product
        self.profile = profile

    def validate_input(self, query, hostname, username):
        if hostname and 'hostname' in query:
            click.echo('Cannot use --hostname with "hostname:" (in query)')
            return False
        elif username and 'username' in query:
            click.echo('Cannot use --username with "username:" (in query)')
            return False
        else:
            return True

    # Build the query based on the product that was chosen
    def base_query(self, *args):
        if self.product == "cbr":
            return cbr.build_query(*args)

        elif self.product == "cbth":
            return args


    # Search based on the product that was chosen
    def process_search(self, cb_conn, base_query, query):
        if self.product == "cbr":
            return cbr.process_search(cb_conn, query, base_query)

        elif self.product == "cbth":
            return cbth.process_search(cb_conn, query, base_query)

    # If defdir or deffiles were given run the appropriate search based on the product
    def nested_process_search(self, criteria, cb_conn, base_query):
        if self.product == "cbr":
            return cbr.nested_process_search(cb_conn, criteria, base_query)

        elif self.product == "cbth":
            return cbth.nested_process_search(cb_conn, criteria, base_query)

    # write the rows of the CSV
    def write_csv(self, output, results, *args):
        for r in results:
            row = [r[0], r[1], r[2], r[3], args[0], args[1]]
            output.writerow(row)

    def get_cbapi_connection(self):
        if self.product == 'cbr':
            if self.profile:
                cb_conn = CbEnterpriseResponseAPI(profile=self.profile)
            else:
                cb_conn = CbEnterpriseResponseAPI()

        elif self.product == 'cbth':
            if self.profile:
                cb_conn = CbThreatHunterAPI(profile=self.profile)
            else:
                cb_conn = CbThreatHunterAPI()
        return cb_conn