Note: The default ITS GitLab runner is a shared resource and is subject to slowdowns during heavy usage.
You can run your own GitLab runner that is dedicated just to your group if you need to avoid processing delays.

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()