1import dshell.core
2from dshell.util import printable_text
3from dshell.output.alertout import AlertOutput
4
5import re
6import sys
7
8class DshellPlugin(dshell.core.ConnectionPlugin):
9
10    def __init__(self):
11        super().__init__(
12            name="search",
13            author="dev195",
14            bpf="tcp or udp",
15            description="Search for patterns in connections",
16            longdescription="""
17Reconstructs streams and searches the content for a user-provided regular
18expression. Requires definition of the --search_expression argument. Additional
19options can be provided to alter behavior.
20            """,
21            output=AlertOutput(label=__name__),
22            optiondict={
23                "expression": {
24                    "help": "Search expression",
25                    "type": str,
26                    "metavar": "REGEX"},
27                "ignorecase": {
28                    "help": "Ignore case when searching",
29                    "action": "store_true"},
30                "invert": {
31                    "help": "Return connections that DO NOT match expression",
32                    "action": "store_true"},
33                "quiet": {
34                    "help": "Do not display matches from this plugin. Useful when chaining plugins.",
35                    "action": "store_true"}
36            })
37
38
39
40    def premodule(self):
41        # make sure the user actually provided an expression to search for
42        if not self.expression:
43            self.error("Must define an expression to search for using --search_expression")
44            sys.exit(1)
45
46        # define the regex flags, based on arguments
47        re_flags = 0
48        if self.ignorecase:
49            re_flags = re_flags | re.IGNORECASE
50
51        # Create the regular expression
52        try:
53            # convert expression to bytes so it can accurately compare to
54            # the connection data (which is also of type bytes)
55            byte_expression = bytes(self.expression, 'utf-8')
56            self.regex = re.compile(byte_expression, re_flags)
57        except Exception as e:
58            self.error("Could not compile regex ({0})".format(e))
59            sys.exit(1)
60
61
62
63    def connection_handler(self, conn):
64        """
65        Go through the data of each connection.
66        If anything is a hit, return the entire connection.
67        """
68
69        match_found = False
70        for blob in conn.blobs:
71            for line in blob.data.splitlines():
72                match = self.regex.search(line)
73                if match and self.invert:
74                    return None
75                elif match and not self.invert:
76                    match_found = True
77                    if not self.quiet:
78                        if blob.sip == conn.sip:
79                            self.write(printable_text(line, False), **conn.info(), dir_arrow="->")
80                        else:
81                            self.write(printable_text(line, False), **conn.info(), dir_arrow="<-")
82                elif self.invert and not match:
83                    if not self.quiet:
84                        self.write(**conn.info())
85                    return conn
86        if match_found:
87            return conn
88
89
90
91if __name__ == "__main__":
92    print(DshellPlugin())
93