1"""
2Only follows connections that match user-provided IP addresses and ports. Is
3generally chained with other plugins.
4"""
5
6import ipaddress
7import sys
8
9import dshell.core
10from dshell.output.alertout import AlertOutput
11
12class DshellPlugin(dshell.core.ConnectionPlugin):
13    def __init__(self, **kwargs):
14        super().__init__(
15            name="track",
16            author="twp,dev195",
17            description="Only follow connections that match user-provided IP addresses and ports",
18            longdescription="""Only follow connections that match user-provided IP addresses
19
20IP addresses can be specified with --track_source and --track_target.
21Multiple IPs can be used with commas (e.g. --track_source=192.168.1.1,127.0.0.1).
22Ports can be included with IP addresses by joining them with a 'p' (e.g. --track_target=192.168.1.1p80,127.0.0.1).
23Ports can be used alone with just a 'p' (e.g. --track_target=p53).
24CIDR notation is okay (e.g. --track_source=196.168.0.0/16).
25
26--track_source : used to limit connections by the IP that initiated the connection (usually the client)
27--trace_target : used to limit connections by the IP that received the connection (usually the server)
28--track_alerts : used to display optional alerts indicating when a connection starts/ends""",
29            bpf="ip or ip6",
30            output=AlertOutput(label=__name__),
31            optiondict={
32                "target": {
33                    "default": [],
34                    "action": "append",
35                    "metavar": "IPpPORT"},
36                "source": {
37                    "default": [],
38                    "action": "append",
39                    "metavar": "IPpPORT"},
40                "alerts": {
41                    "action": "store_true"}
42                }
43            )
44        self.sources = []
45        self.targets = []
46
47    def __split_ips(self, input):
48        """
49        Used to split --track_target and --track_source arguments into
50        list-of-lists used in the connection handler
51        """
52        return_val = []
53        for piece in input.split(','):
54            if 'p' in piece:
55                ip, port = piece.split('p', 1)
56                try:
57                    port = int(port)
58                except ValueError as e:
59                    self.error("Could not parse port number in {!r} - {!s}".format(piece, e))
60                    sys.exit(1)
61                if 0 < port > 65535:
62                    self.error("Could not parse port number in {!r} - must be in valid port range".format(piece))
63                    sys.exit(1)
64            else:
65                ip, port = piece, None
66            if '/' in ip:
67                try:
68                    ip = ipaddress.ip_network(ip)
69                except ValueError as e:
70                    self.error("Could not parse CIDR netrange - {!s}".format(e))
71                    sys.exit(1)
72            elif ip:
73                try:
74                    ip = ipaddress.ip_address(ip)
75                except ValueError as e:
76                    self.error("Could not parse IP address - {!s}".format(e))
77                    sys.exit(1)
78            else:
79                ip = None
80            return_val.append((ip, port))
81        return return_val
82
83    def __check_ips(self, masterip, masterport, checkip, checkport):
84        "Checks IPs and ports for matches against the user-selected values"
85        # masterip, masterport are the values selected by the user
86        # checkip, checkport are the values to be checked against masters
87        ip_okay = False
88        port_okay = False
89
90        if masterip is None:
91            ip_okay = True
92        elif (isinstance(masterip, (ipaddress.IPv4Network, ipaddress.IPv6Network))
93            and checkip in masterip):
94                ip_okay = True
95        elif (isinstance(masterip, (ipaddress.IPv4Address, ipaddress.IPv6Address))
96            and masterip == checkip):
97                ip_okay = True
98
99        if masterport is None:
100            port_okay = True
101        elif masterport == checkport:
102            port_okay = True
103
104        if port_okay and ip_okay:
105            return True
106        else:
107            return False
108
109
110    def premodule(self):
111        if self.target:
112            for tstr in self.target:
113                self.targets.extend(self.__split_ips(tstr))
114        if self.source:
115            for sstr in self.source:
116                self.sources.extend(self.__split_ips(sstr))
117        self.logger.debug("targets: {!s}".format(self.targets))
118        self.logger.debug("sources: {!s}".format(self.sources))
119
120    def connection_handler(self, conn):
121        if self.targets:
122            conn_okay = False
123            for target in self.targets:
124                targetip = target[0]
125                targetport = target[1]
126                serverip = ipaddress.ip_address(conn.serverip)
127                serverport = conn.serverport
128                if self.__check_ips(targetip, targetport, serverip, serverport):
129                    conn_okay = True
130                    break
131            if not conn_okay:
132                return
133
134        if self.sources:
135            conn_okay = False
136            for source in self.sources:
137                sourceip = source[0]
138                sourceport = source[1]
139                clientip = ipaddress.ip_address(conn.clientip)
140                clientport = conn.clientport
141                if self.__check_ips(sourceip, sourceport, clientip, clientport):
142                    conn_okay = True
143                    break
144            if not conn_okay:
145                return
146
147        if self.alerts:
148            self.write("matching connection", **conn.info())
149
150        return conn
151
152if __name__ == "__main__":
153    print(DshellPlugin())
154