sift_labels.py 4.44 KB
Newer Older
1
#!/usr/bin/env python3
2
3
4
5
6
import argparse
import logging
import logging.handlers
import yaml

7
from .core import arg_prompt
8
9
10
11
12
13
14
15
16
from subprocess import check_output as cmd
from subprocess import CalledProcessError


class Sifter():

    def __init__(self, args, config):
        self.__dict__.update(args)
        self.config = config
Rob Carleski's avatar
Rob Carleski committed
17
        self.logger = logging.getLogger(__name__)
18
19
20
21
22
23
24
25

        try:
            cmd([
                config['google']['gam_command'],
                'whatis',
                self.account
            ])
        except CalledProcessError as e:
Rob Carleski's avatar
Rob Carleski committed
26
            self.logger.error(e.output, extra={'entity': self.account})
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
            exit(2)

        # Assemble the list of labels to omit
        omit = [
            'personal',
            'private',
            'personal and private',
            'category_personal',
            'category_social',
            'category_forums',
            'category_updates',
            'category_promotions',
            'chat'
        ]
        if hasattr(self, 'omit_labels'):
            omit = omit + self.omit_labels

        self.omit_labels = list(map((lambda x:  x.lower()), omit))

    def sift(self):
        try:
            user_labels = cmd([
                self.config['google']['gam_command'],
                'user',
                self.account,
                'show',
                'labels'
            ]).split(b'\n\n')[1:-1]
        except CalledProcessError as e:
Rob Carleski's avatar
Rob Carleski committed
56
            self.logger.error(e.output, extra={'entity': self.account})
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

        user_labels = map((lambda label: label.split(b'\n')[0]), user_labels)
        self.clean_labels = []
        for label in user_labels:
            splitLabel = label.split(b'/')
            try:
                for part in splitLabel:
                    assert part.decode('utf-8').lower() not in self.omit_labels
                self.clean_labels.append(label)
            except AssertionError:
                pass

    def print_sifted_labels(self):
        for _ in self.clean_labels:
            print(_.decode('utf-8'))


def main():

    helptext = '''examples:
    google-sift-labels -a mantony -o 'alesia_notes'
    google-sift-labels --account mantony --omit_labels 'alesia_notes'
    '''

    # Collect arguments for the script.
    parser = argparse.ArgumentParser(
        description='Reports on user labels in gmail.',
        epilog=helptext,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        '--account',
        '-a',
        help="The user to check files for.",
    )
    parser.add_argument(
        '--omit_labels',
        '-o',
        help='A list of additional labels to omit from the report.',
        nargs='+',
        default=[]
    )
    parser.add_argument(
        '--config',
        '-c',
        help='The CAK configuration to use.',
        default='/etc/collab-admin-kit.yml'
    )
    args = parser.parse_args()

107
108
109
110
111
112
113
114
    # Argument prompt fallback
    if not args.account:
        args.account = arg_prompt(
            'Sifter target account'
        )
        args.omit_labels = arg_prompt(
            '(optional) Omit labels with the following names',
            default=''
Rob Carleski's avatar
Rob Carleski committed
115
        ).split(',')
116
117
118
119
120
121
122
123
124
125
126
127

    # Argument sanity check
    _required = ['account']
    for r in _required:
        try:
            if isinstance(r, str):
                assert getattr(args, r)
            if isinstance(r, tuple):
                assert getattr(args, r[0]) or getattr(args, r[1])
        except AssertionError:
            print('Missing required argument {}'.format(r))

128
129
130
131
132
    # Open the CAK Config
    with open(args.config) as stream:
        config = yaml.load(stream, Loader=yaml.BaseLoader)

    # Get the root logger and set the debug level
Rob Carleski's avatar
Rob Carleski committed
133
    logger = logging.getLogger(__name__)
134
135
136
137
138
139
140
141
142
143
144
145
146
    logger.setLevel(logging.DEBUG)

    # Create a syslog handler, set format, and associate.
    sh = logging.handlers.SysLogHandler(
        address='/dev/log',
        facility=config['general']['log_facility']
    )
    formatter = logging.Formatter(config['general']['log_format'])
    sh.setFormatter(formatter)
    logger.addHandler(sh)

    # Create a console handler, set format, and associate.
    ch = logging.StreamHandler()
147
148
149
150
    formatter = logging.Formatter(
        config['general']['console_format'],
        config['general']['date_format']
    )
151
152
153
    ch.setFormatter(formatter)
    logger.addHandler(ch)

154
    sifter = Sifter(vars(args), config)
155
156
157
158
159
160
    sifter.sift()
    sifter.print_sifted_labels()


if __name__ == '__main__':
    main()