← Back to branch summary

~ubuntuone-hackers/conn-check/trunk

« back to all changes in this revision

Viewing changes to conn_check/utils/firewall_rules.py

[r=wesmason] Add firewall rule yaml output options

Show diffs side-by-side

added added

removed removed

Lines of Context:
conn_check/utils/firewall_rules.py
 
1
import sys
 
2
import socket
 
3
import yaml
 
4
 
 
5
from ..main import Command, parse_version_arg
 
6
 
 
7
 
 
8
class FirewallRulesOutput(object):
 
9
    """Outputs a set of YAML firewall rules matching checks."""
 
10
 
 
11
    def __init__(self, output):
 
12
        self.output = output
 
13
        self.output_data = {}
 
14
        self.fqdn = socket.getfqdn()
 
15
 
 
16
    def notify_skip(self, name):
 
17
        """Passes skips.
 
18
 
 
19
        Should pass everything when using the skipping_strategy, directly to
 
20
        write().
 
21
        """
 
22
 
 
23
        self.write(name)
 
24
 
 
25
    def write(self, data):
 
26
        """Filters out non-TCP/UDP checks.
 
27
 
 
28
        Stores host/port/proto info for output later as YAML.
 
29
        """
 
30
 
 
31
        # We only need TCP/UDP checks
 
32
        if not any(x in data for x in ('tcp', 'udp')):
 
33
            return
 
34
 
 
35
        # Here we take the list of colon separated values in reverse order, so
 
36
        # we should get just the host/port/proto for the check without the
 
37
        # specific prefix (e.g. memcache, http)
 
38
        port, host, protocol = reversed(data.split(':')[-3:])
 
39
        protocol = protocol.strip()
 
40
 
 
41
        key = "{}:{}".format(host, protocol)
 
42
        if key not in self.output_data:
 
43
            self.output_data[key] = {
 
44
                'from_host': self.fqdn,
 
45
                'to_host': host,
 
46
                'ports': [],
 
47
                'protocol': protocol,
 
48
            }
 
49
 
 
50
        port = int(port)
 
51
        if port not in self.output_data[key]['ports']:
 
52
            self.output_data[key]['ports'].append(port)
 
53
 
 
54
    def flush(self):
 
55
        """Outputs our structured egress firewall info as YAML."""
 
56
 
 
57
        self.output.write(yaml.dump({'egress': self.output_data.values()}))
 
58
 
 
59
 
 
60
class FirewallExportCommand(Command):
 
61
    """CLI command runner for conn-check-export-fw"""
 
62
 
 
63
    def wrap_output(self, output):
 
64
        """Wraps output stream.
 
65
 
 
66
        Override some options in order to just output fw rules without
 
67
        performing checks.
 
68
        """
 
69
        # We don't want to actually perform the checks
 
70
        self.options.dry_run = True
 
71
        self.options.buffer_output = False
 
72
        self.options.show_duration = False
 
73
 
 
74
        super(FirewallExportCommand, self).wrap_output(output)
 
75
 
 
76
        self.output = FirewallRulesOutput(self.output)
 
77
        self.results.output = self.output
 
78
 
 
79
 
 
80
def run(*args):
 
81
    if parse_version_arg():
 
82
        return 0
 
83
 
 
84
    cmd = FirewallExportCommand(args)
 
85
    return cmd.run()
 
86
 
 
87
 
 
88
def main():
 
89
    sys.exit(run(*sys.argv[1:]))
 
90
 
 
91
 
 
92
if __name__ == '__main__':
 
93
    main()