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