transfer_mail.py 9.79 KB
Newer Older
1
#!/usr/bin/env python3
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import argparse
import logging
import logging.handlers
import os
import re
import yaml

from hashlib import md5
from shutil import rmtree
from subprocess import CalledProcessError
from subprocess import check_output as cmd

devnull = open(os.devnull, 'w')


class Mover:

    def __init__(self, args, config):
        self.__dict__.update(args)
        self.config = config
Rob Carleski's avatar
Rob Carleski committed
22
        self.logger = logging.getLogger(__name__)
23
24
25
26
27
28
29
30
31
32
33
34
35

        try:
            cmd([
                self.config['google']['gam_command'],
                'whatis',
                self.current_owner
            ])
            cmd([
                self.config['google']['gam_command'],
                'whatis',
                self.new_owner
            ])
        except CalledProcessError as e:
Rob Carleski's avatar
Rob Carleski committed
36
            self.logger.error(e.output, extra={'entity': self.current_owner})
37
38
39
40
41
42
43
44
45
            exit(2)

        try:
            with open(self.label_file) as labelFile:
                self.labels = []
                labels = labelFile.read().splitlines()
                for label in filter(None, labels):
                    self.labels.append(label)
        except OSError as e:
Rob Carleski's avatar
Rob Carleski committed
46
            self.logger.error(
47
48
49
50
51
52
53
54
                e.strerror,
                extra={'entity': self.current_owner}
            )
            exit(2)
        except AttributeError as e:
            pass

        if len(self.labels) < 1:
Rob Carleski's avatar
Rob Carleski committed
55
            self.logger.error(
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
                'No labels found to transfer.',
                extra={'entity': self.current_owner}
            )
            exit(2)

        # Create list of labels to work with. Entries in this list
        # are tuples of the label search name and filesystem-safe download
        # directory name. This gives us a way to have not only the normal
        # label name to push back up, but also something to search properly
        # with, and something we can store mail on disk with.
        self.transfer = {}
        for label in self.labels:
            self.transfer[label] = {}
            self.transfer[label]['original'] = label
            self.transfer[label]['searchable'] = re.sub(
                r"[\^\|\&/)(\s]",
                '-',
                label
            )
75
            if self.label_prefix:
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
                self.transfer[label]['destination'] = '{}-{}'.format(
                    self.label_prefix,
                    label
                )
            else:
                self.transfer[label]['destination'] = label

            self.transfer[label]['disk_hash'] = md5(
                label.encode('UTF-8')
            ).hexdigest()[:8]
            self.transfer[label]['dest_searchable'] = re.sub(
                r"[\^\|\&\/\)\(\s]",
                '-',
                self.transfer[label]['destination']
            )

        # Create a probably-unique hash to use for collision avoidance
        ownerhash = md5('{}{}'.format(
            self.current_owner,
            self.new_owner
            ).encode('UTF-8')
        ).hexdigest()[:8]

Rob Carleski's avatar
Rob Carleski committed
99
        if self.email_directory:
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
            pass
        else:
            self.email_directory = '{}/{}/{}/'.format(
                self.config['general']['data_dir'],
                'mail',
                ownerhash
            )

        if not os.path.exists(self.email_directory):
            os.makedirs(self.email_directory)

        if '@' not in self.current_owner:
            self.current_owner += '@' + config['google']['domain']
        self.current_owner = self.current_owner

        if '@' not in self.new_owner:
            self.new_owner += '@' + config['google']['domain']
        self.new_owner = self.new_owner

    def transfer_mail(self):
        for _ in self.transfer:
            self.process_label(self.transfer[_])

    def process_label(self, label):
        origin_count = self.count_label(self.current_owner, label)
        if origin_count > 0:
            self.pull_label(label)
            self.push_label(label)

            destination_count = self.count_label(self.new_owner, label)
            if origin_count > destination_count:
Rob Carleski's avatar
Rob Carleski committed
131
                self.logger.warning(
132
133
134
135
136
                    'Message count mismatch. Retrying transfer.',
                    extra={'entity': self.current_owner}
                )
                self.process_label(label)
            else:
Rob Carleski's avatar
Rob Carleski committed
137
                self.logger.info(
138
139
140
141
                    'Message counts match between mailboxes',
                    extra={'entity': self.current_owner}
                )
        else:
Rob Carleski's avatar
Rob Carleski committed
142
            self.logger.info(
143
144
145
146
147
                    'No messages found for {}'.format(label['original']),
                    extra={'entity': self.current_owner}
            )

    def pull_label(self, label):
Rob Carleski's avatar
Rob Carleski committed
148
        self.logger.info(
Rob Carleski's avatar
Rob Carleski committed
149
            'Pulling messages in label {}'.format(label['original']),
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
            extra={'entity': self.current_owner}
        )
        try:
            cmd([
                self.config['google']['gyb_command'],
                '--email',
                self.current_owner,
                '--action',
                'backup',
                '--search',
                'label:{}'.format(label['searchable']),
                '--local-folder',
                self.email_directory + label['disk_hash']
                ], stderr=devnull)
        except CalledProcessError as e:
Rob Carleski's avatar
Rob Carleski committed
165
            self.logger.warning(e.output, extra={'entity': self.current_owner})
166
167

    def push_label(self, label):
Rob Carleski's avatar
Rob Carleski committed
168
        self.logger.info(
Rob Carleski's avatar
Rob Carleski committed
169
            'Pushing messages in label {}'.format(label['original']),
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
            extra={'entity': self.current_owner}
        )
        try:
            cmd([
                self.config['google']['gyb_command'],
                '--email',
                self.new_owner,
                '--action',
                'restore',
                '--strip-labels',
                '--local-folder',
                self.email_directory + label['disk_hash'],
                '--label-restored',
                '{}'.format(label['destination'])
            ])
        except CalledProcessError as e:
Rob Carleski's avatar
Rob Carleski committed
186
            self.logger.warning(e.output, extra={'entity': self.current_owner})
187
188
189
190
191
192
193
194
195

    def count_label(self, user, label):
        if user == self.current_owner:
            search_label = label['searchable']
            location = 'origin'
        else:
            search_label = label['dest_searchable']
            location = 'destination'

Rob Carleski's avatar
Rob Carleski committed
196
        self.logger.info(
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
            'Counting messages in {} at {}'.format(
                label['original'],
                location
            ),
            extra={'entity': self.current_owner}
        )
        count = cmd([
            self.config['google']['gyb_command'],
            '--email',
            user,
            '--action',
            'count',
            '--search',
            'label:{}'.format(
                search_label
            )]
        )
        count = int(count.split(b',')[1])

Rob Carleski's avatar
Rob Carleski committed
216
        self.logger.info(
217
218
219
220
221
222
223
224
225
226
227
228
229
            'Found {} messages in {} at {}'.format(
                count,
                label['original'],
                location
            ),
            extra={'entity': self.current_owner}
        )
        return count

    def cleanup(self):
        try:
            rmtree(self.email_directory)
        except OSError:
Rob Carleski's avatar
Rob Carleski committed
230
            self.logger.warning(
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
                'Unable to delete email archive folder at {}{}'.format(
                    self.config['general']['data_dir'],
                    self.email_directory
                )
            )


def main():

    helptext = '''examples:
    google-transfer-mail -o dgrohl -n cweathrs -f '/tmp/labels' -p 'from dave'
    google-transfer-mail --current_owner dgrohl --new_owner cweathrs
        --label_prefix 'from_dave' --label_list foofighters qotsa nirvana
    '''

    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description='Transfers email between two gsuite accounts',
        epilog=helptext,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        '--current_owner',
        '-o',
        help='The original owner of the email message(s)',
        required=True
    )
    parser.add_argument(
        '--new_owner',
        '-n',
        help='The new owner of the email message(s)',
        required=True
    )
    labels = parser.add_mutually_exclusive_group(required=True)
    labels.add_argument(
        '--label_file',
        '-f',
        help='A newline-delimited list of labels to be transferred.',
    )
    labels.add_argument(
        '--labels',
        '-l',
        help='A list of labels to be transferred',
        nargs='+'
    )
    parser.add_argument(
        '--label_prefix',
        '-p',
        help='A prefix to append to existing label names',
    )
    parser.add_argument(
        '--email_directory',
        '-d',
        help='Subdirectory of GAK data directory to store mail in.'
    )
    parser.add_argument(
        '-c',
        '--config',
        help="The GAK config to use.",
        default='/etc/collab-admin-kit.yml'
    )
    args = parser.parse_args()

    # 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
299
    logger = logging.getLogger(__name__)
300
301
302
303
304
    logger.setLevel(logging.DEBUG)

    # Create a syslog handler, set format, and associate.
    sh = logging.handlers.SysLogHandler(
        address='/dev/log',
Rob Carleski's avatar
Rob Carleski committed
305
        facility=config['general']['log_facility']
306
    )
Rob Carleski's avatar
Rob Carleski committed
307
    formatter = logging.Formatter(config['general']['log_format'])
308
309
310
311
312
    sh.setFormatter(formatter)
    logger.addHandler(sh)

    # Create a console handler, set format, and associate.
    ch = logging.StreamHandler()
Rob Carleski's avatar
Rob Carleski committed
313
    formatter = logging.Formatter(config['general']['console_format'])
314
315
316
    ch.setFormatter(formatter)
    logger.addHandler(ch)

317
    mover = Mover(vars(args), config)
318
319
320
321
322
323
    mover.transfer_mail()
    mover.cleanup()


if __name__ == '__main__':
    main()