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.

transfer_drive.py 5.91 KB
Newer Older
1
#!/usr/bin/env python3
2
import argparse
3
import csv
4
5
6
7
8
import logging
import logging.handlers
import os
import yaml

9
from .core import arg_prompt
10
11
12
13
14
15
16
17
18
19
20
from subprocess import check_output as cmd
from subprocess import CalledProcessError

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
21
        self.logger = logging.getLogger(__name__)
22
23
24
25
26

        try:
            cmd([
                self.config['google']['gam_command'],
                'whatis',
27
                self.owner
28
29
            ])
        except CalledProcessError as e:
30
            self.logger.error(e.output, extra={'entity': self.owner})
31
32
33
34
35
            exit(2)

        if os.path.isfile(self.id_file):
            pass
        else:
Rob Carleski's avatar
Rob Carleski committed
36
            self.logger.error(
37
                'Unable to find given ID file',
38
                extra={'entity': self.owner}
39
40
41
42
            )
            exit(2)

    def transfer_files(self):
43
44
45
        with open(self.id_file) as csvfile:
            reader = csv.DictReader(csvfile)
            for item in reader:
46
47
48
49
                try:
                    cmd([
                        self.config['google']['gam_command'],
                        'user',
50
                        item['owner'],
51
52
                        'update',
                        'drivefileacl',
53
54
                        item['id'],
                        self.owner,
55
56
57
                        'role',
                        'owner'
                    ], stderr=devnull)
58
                except CalledProcessError:
59
60
61
62
                    try:
                        cmd([
                            self.config['google']['gam_command'],
                            'user',
63
                            item['owner'],
64
65
                            'add',
                            'drivefileacl',
66
                            item['id'],
67
                            'user',
68
                            self.owner,
69
70
                            'role',
                            'reader'
Rob Carleski's avatar
Rob Carleski committed
71
                        ])
72
73
74
                        cmd([
                            self.config['google']['gam_command'],
                            'user',
75
                            item['owner'],
76
77
                            'update',
                            'drivefileacl',
78
79
                            item['id'],
                            self.owner,
80
81
                            'role',
                            'owner'
Rob Carleski's avatar
Rob Carleski committed
82
                        ])
83
                    except CalledProcessError as e:
Rob Carleski's avatar
Rob Carleski committed
84
                        self.logger.warning(
85
                            e.output,
86
                            extra={'entity': item['owner']}
87
88
89
                        )
                        return None
                if hasattr(self, 'strip_access'):
90
91
                    if self.strip_access:
                        self.strip_permissions(item)
92

93
    def strip_permissions(self, item):
94
95
96
97
        try:
            cmd([
                self.config['google']['gam_command'],
                'user',
98
                self.owner,
99
100
                'delete',
                'drivefileacl',
101
102
                item['id'],
                item['owner']
Rob Carleski's avatar
Rob Carleski committed
103
                ])
104
        except CalledProcessError as e:
105
            self.logger.warning(e.output, extra={'entity': item['owner']})
106
107
108
109
110


def main():

    helptext = '''examples:
111
112
113
    google-transer-drive -o obloom -f '/tmp/obloom-files' -s
    google-transfer-drive --owner obloom --id_file '/tmp/obloom-files'
        --strip_access
114
115
116
117
118
119
120
121
    '''

    parser = argparse.ArgumentParser(
        description='Transfers files between users in a GSuite domain',
        epilog=helptext,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
122
        '--owner',
123
124
125
126
127
128
        '-o',
        help="The new owner of the files being transferred.",
    )
    parser.add_argument(
        '--id_file',
        '-f',
129
        help='transfer_files.csv from google-sift-drive.',
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
    )
    parser.add_argument(
        '--strip_access',
        '-s',
        help='Strips the original owners access from file(s).',
        action='store_true'
    )
    parser.add_argument(
        '--config',
        '-c',
        help='The CAK config file to use.',
        default='/etc/collab-admin-kit.yml'
    )
    args = parser.parse_args()

145
146
147
    # Argument prompt fallback
    if not args.owner:
        args.owner = arg_prompt(
Rob Carleski's avatar
Rob Carleski committed
148
            'Uniqname of new file/folder owner'
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
        )
        args.id_file = arg_prompt(
            'Full path to ID file generated from google-sift-drive'
        )
        args.strip_access = arg_prompt(
            '(optional) Strip access from original owner [y/N]',
            default='n'
        ).lower()

    # Argument sanity check
    _required = ['owner', 'id_file']
    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))

169
170
171
172
173
    # 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
174
    logger = logging.getLogger(__name__)
175
176
177
178
179
    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
180
        facility=config['general']['log_facility']
181
    )
Rob Carleski's avatar
Rob Carleski committed
182
    formatter = logging.Formatter(config['general']['log_format'])
183
184
185
186
187
    sh.setFormatter(formatter)
    logger.addHandler(sh)

    # Create a console handler, set format, and associate.
    ch = logging.StreamHandler()
188
189
190
191
    formatter = logging.Formatter(
        config['general']['console_format'],
        config['general']['date_format']
    )
192
193
194
    ch.setFormatter(formatter)
    logger.addHandler(ch)

195
    mover = Mover(vars(args), config)
196
197
198
199
200
    mover.transfer_files()


if __name__ == '__main__':
    main()