transfer_drive.py

#!/usr/bin/env python3
import argparse
import csv
import logging
import logging.handlers
import os
import yaml

from .core import arg_prompt
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
        self.logger = logging.getLogger(__name__)

        try:
            cmd([
                self.config['google']['gam_command'],
                'whatis',
                self.owner
            ])
        except CalledProcessError as e:
            self.logger.error(e.output, extra={'entity': self.owner})
            exit(2)

        if os.path.isfile(self.id_file):
            pass
        else:
            self.logger.error(
                'Unable to find given ID file',
                extra={'entity': self.owner}
            )
            exit(2)

    def transfer_files(self):
        with open(self.id_file) as csvfile:
            reader = csv.DictReader(csvfile)
            for item in reader:
                try:
                    cmd([
                        self.config['google']['gam_command'],
                        'user',
                        item['owner'],
                        'update',
                        'drivefileacl',
                        item['id'],
                        self.owner,
                        'role',
                        'owner'
                    ], stderr=devnull)
                except CalledProcessError:
                    try:
                        cmd([
                            self.config['google']['gam_command'],
                            'user',
                            item['owner'],
                            'add',
                            'drivefileacl',
                            item['id'],
                            'user',
                            self.owner,
                            'role',
                            'reader'
                        ])
                        cmd([
                            self.config['google']['gam_command'],
                            'user',
                            item['owner'],
                            'update',
                            'drivefileacl',
                            item['id'],
                            self.owner,
                            'role',
                            'owner'
                        ])
                    except CalledProcessError as e:
                        self.logger.warning(
                            e.output,
                            extra={'entity': item['owner']}
                        )
                        return None
                if hasattr(self, 'strip_access'):
                    if self.strip_access:
                        self.strip_permissions(item)

    def strip_permissions(self, item):
        try:
            cmd([
                self.config['google']['gam_command'],
                'user',
                self.owner,
                'delete',
                'drivefileacl',
                item['id'],
                item['owner']
                ])
        except CalledProcessError as e:
            self.logger.warning(e.output, extra={'entity': item['owner']})


def main():

    helptext = '''examples:
    google-transer-drive -o obloom -f '/tmp/obloom-files' -s
    google-transfer-drive --owner obloom --id_file '/tmp/obloom-files'
        --strip_access
    '''

    parser = argparse.ArgumentParser(
        description='Transfers files between users in a GSuite domain',
        epilog=helptext,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        '--owner',
        '-o',
        help="The new owner of the files being transferred.",
    )
    parser.add_argument(
        '--id_file',
        '-f',
        help='transfer_files.csv from google-sift-drive.',
    )
    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()

    # Argument prompt fallback
    if not args.owner:
        args.owner = arg_prompt(
            'Uniqname of new file/folder owner'
        )
        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))

    # 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
    logger = logging.getLogger(__name__)
    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()
    formatter = logging.Formatter(
        config['general']['console_format'],
        config['general']['date_format']
    )
    ch.setFormatter(formatter)
    logger.addHandler(ch)

    mover = Mover(vars(args), config)
    mover.transfer_files()


if __name__ == '__main__':
    main()