Commit 72f63140 authored by Rob Carleski's avatar Rob Carleski 🇮🇸
Browse files

Use csv-style input file for transfer_drive

parent 4983b862
#!/usr/bin/env python3 #!/usr/bin/env python3
import argparse import argparse
import csv
import logging import logging
import logging.handlers import logging.handlers
import os import os
...@@ -22,15 +23,10 @@ class Mover: ...@@ -22,15 +23,10 @@ class Mover:
cmd([ cmd([
self.config['google']['gam_command'], self.config['google']['gam_command'],
'whatis', 'whatis',
self.current_owner self.owner
])
cmd([
self.config['google']['gam_command'],
'whatis',
self.new_owner
]) ])
except CalledProcessError as e: except CalledProcessError as e:
self.logger.error(e.output, extra={'entity': self.current_owner}) self.logger.error(e.output, extra={'entity': self.owner})
exit(2) exit(2)
if os.path.isfile(self.id_file): if os.path.isfile(self.id_file):
...@@ -38,23 +34,23 @@ class Mover: ...@@ -38,23 +34,23 @@ class Mover:
else: else:
self.logger.error( self.logger.error(
'Unable to find given ID file', 'Unable to find given ID file',
extra={'entity': self.current_owner} extra={'entity': self.owner}
) )
exit(2) exit(2)
def transfer_files(self): def transfer_files(self):
with open(self.id_file) as stream: with open(self.id_file) as csvfile:
for file_id in stream: reader = csv.DictReader(csvfile)
file_id = file_id.strip() for item in reader:
try: try:
cmd([ cmd([
self.config['google']['gam_command'], self.config['google']['gam_command'],
'user', 'user',
self.current_owner, item['owner'],
'update', 'update',
'drivefileacl', 'drivefileacl',
file_id, item['id'],
self.new_owner, self.owner,
'role', 'role',
'owner' 'owner'
], stderr=devnull) ], stderr=devnull)
...@@ -63,48 +59,49 @@ class Mover: ...@@ -63,48 +59,49 @@ class Mover:
cmd([ cmd([
self.config['google']['gam_command'], self.config['google']['gam_command'],
'user', 'user',
self.current_owner, item['owner'],
'add', 'add',
'drivefileacl', 'drivefileacl',
file_id, item['id'],
'user', 'user',
self.new_owner, self.owner,
'role', 'role',
'reader' 'reader'
], stderr=devnull) ], stderr=devnull)
cmd([ cmd([
self.config['google']['gam_command'], self.config['google']['gam_command'],
'user', 'user',
self.current_owner, item['owner'],
'update', 'update',
'drivefileacl', 'drivefileacl',
file_id, item['id'],
self.new_owner, self.owner,
'role', 'role',
'owner' 'owner'
], stderr=devnull) ], stderr=devnull)
except CalledProcessError as e: except CalledProcessError as e:
self.logger.warning( self.logger.warning(
e.output, e.output,
extra={'entity': self.current_owner} extra={'entity': item['owner']}
) )
return None return None
if hasattr(self, 'strip_access'): if hasattr(self, 'strip_access'):
self.strip_permissions(file_id) if self.strip_access:
self.strip_permissions(item)
def strip_permissions(self, file_id): def strip_permissions(self, item):
try: try:
cmd([ cmd([
self.config['google']['gam_command'], self.config['google']['gam_command'],
'user', 'user',
self.new_owner, self.owner,
'delete', 'delete',
'drivefileacl', 'drivefileacl',
file_id, item['id'],
self.current_owner item['owner']
], stderr=devnull) ], stderr=devnull)
except CalledProcessError as e: except CalledProcessError as e:
self.logger.warning(e.output, extra={'entity': self.current_owner}) self.logger.warning(e.output, extra={'entity': item['owner']})
def main(): def main():
...@@ -115,28 +112,21 @@ def main(): ...@@ -115,28 +112,21 @@ def main():
'/tmp/obloom-files' --strip_access '/tmp/obloom-files' --strip_access
''' '''
# Kick up an argument parser and assign values into the env object.
parser = argparse.ArgumentParser( parser = argparse.ArgumentParser(
description='Transfers files between users in a GSuite domain', description='Transfers files between users in a GSuite domain',
epilog=helptext, epilog=helptext,
formatter_class=argparse.RawDescriptionHelpFormatter formatter_class=argparse.RawDescriptionHelpFormatter
) )
parser.add_argument( parser.add_argument(
'--current_owner', '--owner',
'-o', '-o',
help="The current owner of the files being transferred.",
required=True
)
parser.add_argument(
'--new_owner',
'-n',
help="The new owner of the files being transferred.", help="The new owner of the files being transferred.",
required=True required=True
) )
parser.add_argument( parser.add_argument(
'--id_file', '--id_file',
'-f', '-f',
help='The file containing a newline-separated list of file IDs.', help='transfer_files.csv from google-sift-drive.',
required=True required=True
) )
parser.add_argument( parser.add_argument(
......
afkdjfglksdjglksjd234_sdgfsdlkfgj owner,id
13584903634l5mn2345lkj540fdb9id-c90v testuser@domain.com,afkdjfglksdjglksjd234_sdgfsdlkfgj
1234h3r-089guda-0gf98u203495erfgdfg testuser@domain.com,13584903634l5mn2345lkj540fdb9id-c90v
3-905483405fdgposdgtlk2j536l4kj5tgdf testuser@domain.com,1234h3r-089guda-0gf98u203495erfgdfg
2k435hj354632848ef09g8dcx09f871235iujh testuser@domain.com,3-905483405fdgposdgtlk2j536l4kj5tgdf
8437659er87g-0f89d7gs9087g89d072341 testuser@domain.com,2k435hj354632848ef09g8dcx09f871235iujh
testuser@domain.com,8437659er87g-0f89d7gs9087g89d072341
...@@ -9,8 +9,7 @@ config = data_dir + 'test_cak_config.yml' ...@@ -9,8 +9,7 @@ config = data_dir + 'test_cak_config.yml'
def test_transfer(): def test_transfer():
args = { args = {
'current_owner': 'testuser', 'owner': 'testuser',
'new_owner': 'testuser',
'id_file': data_dir + '/mock_gam/test_transfer_drive.txt' 'id_file': data_dir + '/mock_gam/test_transfer_drive.txt'
} }
mover = transfer_drive.Mover( mover = transfer_drive.Mover(
...@@ -22,8 +21,7 @@ def test_transfer(): ...@@ -22,8 +21,7 @@ def test_transfer():
def test_transfer_with_strip(): def test_transfer_with_strip():
args = { args = {
'current_owner': 'testuser', 'owner': 'testuser',
'new_owner': 'testuser',
'id_file': data_dir + '/mock_gam/test_transfer_drive.txt', 'id_file': data_dir + '/mock_gam/test_transfer_drive.txt',
'strip_permissions': True 'strip_permissions': True
} }
...@@ -34,25 +32,10 @@ def test_transfer_with_strip(): ...@@ -34,25 +32,10 @@ def test_transfer_with_strip():
mover.transfer_files() mover.transfer_files()
def test_transfer_with_current_owner(): def test_transfer_with_bad_new_owner():
with pytest.raises(SystemExit): with pytest.raises(SystemExit):
args = { args = {
'current_owner': 'notauser', 'owner': 'notauser',
'new_owner': 'testuser',
'id_file': data_dir + '/mock_gam/test_transfer_drive.txt'
}
mover = transfer_drive.Mover(
args,
yaml.load(open(config), Loader=yaml.BaseLoader)
)
mover.transfer_files()
def test_transfer_with_new_owner():
with pytest.raises(SystemExit):
args = {
'current_owner': 'testuser',
'new_owner': 'notauser',
'id_file': data_dir + '/mock_gam/test_transfer_drive.txt' 'id_file': data_dir + '/mock_gam/test_transfer_drive.txt'
} }
mover = transfer_drive.Mover( mover = transfer_drive.Mover(
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment