#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import hashlib
import io
import json
import math
import mmap
import ntpath
import os
import sys
from mimetypes import MimeTypes
from googleapiclient.http import MediaIoBaseDownload
from googleapiclient.http import MediaIoBaseUpload
from size_formatting import formatter
from tabulate import tabulate
from tqdm import tqdm
import encoder
import file_parts
from api import *
from api import GoogleAPI
from custom_exceptions import PythonVersionError, NoClientSecretError, Error
try:
from urllib.request import pathname2url
except ImportError:
Error.formatter(PythonVersionError, ".".join(str(v) for v in sys.version_info[:2]))
if not os.path.exists(os.path.join(os.getcwd() + "/client_secret.json")):
Error.formatter(NoClientSecretError)
DOWNLOADS_FOLDER = "downloads"
TEMP_FOLDER = "tmp"
USE_MULTITHREADED_UPLOADS = True
MAX_DOC_LENGTH = 1000000
MAX_RAM_MB = 1024
MAX_WORKERS_ALLOWED = 10
CHUNK_READ_LENGTH_BYTES = 750000
class UDS:
def __init__(self):
self.api = GoogleAPI()
def delete_file(self, id, name=None, mode_=None):
"""Deletes a given file
Use the Google Drive API to delete a file given its ID.
Args:
id (str): ID of the file
name (str): Name of the file
:param id:
:param name:
:param mode_:
"""
try:
self.api.delete_file(id)
if name is not None:
# If Alpha commands are used, this displays the name
print("Deleted %s" % name)
else:
# If UDS commands are used, this displays the ID
print("Deleted %s" % id)
except IOError:
if mode_ != "quiet":
print("File \"%s\" was not a UDS file" % GoogleAPI.ERROR_OUTPUT)
def build_file(self, parent_id):
"""Download a uds file
This will fetch the Docs one by one, concatenating them
to a local base64 file. The file will then be converted
from base64 to the appropriate mime-type.
Args:
parent_id (str): The ID of the containing folder
:return:
:param parent_id:
"""
items = self.api.recursive_list_folder(parent_id)
folder = self.api.get_file(parent_id)
if not items:
print('No parts found.')
return
# Fix part as int
for item in items:
item['properties']['part'] = int(item['properties']['part'])
items.sort(key=lambda x: x['properties']['part'], reverse=False)
f = open("%s/%s" % (get_downloads_folder(), folder['name']), "a+b")
progress_bar_chunks = tqdm(total=len(items),
unit='chunks', dynamic_ncols=True, position=0)
progress_bar_speed = tqdm(total=len(items) * CHUNK_READ_LENGTH_BYTES, unit_scale=1,
unit='B', dynamic_ncols=True, position=1)
for item in items:
encoded_part = self.download_part(item['id'])
# Decode
decoded_part = encoder.decode(encoded_part)
progress_bar_chunks.update(1)
progress_bar_speed.update(CHUNK_READ_LENGTH_BYTES)
# Append decoded part to file
f.write(decoded_part)
print()
file_hash = self.hash_file(f.name)
f.close()
original_hash = folder.get("md5Checksum")
if original_hash is not None and file_hash != original_hash:
print("Failed to verify hash\nDownloaded file had hash {} compared to original {}".format(
file_hash, original_hash))
os.remove(f.name)
def download_part(self, part_id):
"""
:param part_id:
:return:
"""
request = self.api.export_media(part_id)
fh = io.BytesIO()
downloader = MediaIoBaseDownload(fh, request)
done = False
while done is False:
status, done = downloader.next_chunk()
return fh.getvalue()
def upload_chunked_part(self, chunk, api=None):
"""Upload a chunked part to drive and return the size of the chunk
:param chunk:
:param api:
:return:
"""
if not api:
api = self.api
with open(chunk.path) as fd:
mm = mmap.mmap(fd.fileno(), 0, access=mmap.ACCESS_READ)
chunk_bytes = mm[chunk.range_start:chunk.range_end]
encoded_chunk = encoder.encode(chunk_bytes)
file_metadata = {
'name': chunk.media.name + str(chunk.part),
'mimeType': 'application/vnd.google-apps.document',
'parents': [chunk.parent],
'properties': {
'part': str(chunk.part)
}
}
mediaio_file = MediaIoBaseUpload(io.StringIO(encoded_chunk),
mimetype='text/plain')
self.api.upload_single_file(mediaio_file, file_metadata)
return len(chunk_bytes)
def do_chunked_upload(self, path):
"""
:rtype: object
:param path:
"""
# Prepare media file
size = os.stat(path).st_size
file_hash = self.hash_file(path)
encoded_size = size * (4 / 3)
root = self.api.get_base_folder()['id']
media = file_parts.UDSFile(ntpath.basename(path), None,
MimeTypes().guess_type(pathname2url(path))[0],
formatter(size), formatter(encoded_size), parents=[root], size_numeric=size,
md5=file_hash)
parent = self.api.create_media_folder(media)
# Should be the same
no_docs = math.ceil(encoded_size / MAX_DOC_LENGTH)
# Append all chunks to chunk list
chunk_list = [file_parts.Chunk(path, i, size, media=media, parent=parent['id']) for i in range(no_docs)]
total = 0
total_chunks = no_docs
progress_bar_chunks = tqdm(total=total_chunks,
unit='chunks', dynamic_ncols=True, position=0)
progress_bar_speed = tqdm(total=total_chunks * CHUNK_READ_LENGTH_BYTES, unit_scale=1,
unit='B', dynamic_ncols=True, position=1)
for chunk in chunk_list:
total += 1
self.upload_chunked_part(chunk)
progress_bar_speed.update(CHUNK_READ_LENGTH_BYTES)
progress_bar_chunks.update(1)
#
# Concurrently execute chunk upload and report back when done.
# with concurrent.futures.ProcessPoolExecutor(max_workers=MAX_WORKERS_ALLOWED) as executor:
# for file in executor.map(ext_upload_chunked_part, chunk_list):
# total = total + file
# elapsed_time = round(time.time() - start_time, 2)
# current_speed = round(
# (total) / (elapsed_time * 1024 * 1024), 2)
# progress_bar("Uploading %s at %sMB/s" %
# (media.name, current_speed), total, size)
#
# Print new file output
table = [[media.name, media.size, media.encoded_size, parent['id']]]
print()
print("\n" + tabulate(table, headers=[
'Name', 'Size', 'Encoded', 'ID', ]))
def convert_file(self, file_id):
# Get file metadata
metadata = service.files().get(fileId=file_id, fields="name").execute()
# Download the file and then call do_upload() on it
request = service.files().get_media(fileId=file_id)
path = "%s/%s" % (get_downloads_folder(), metadata['name'])
fh = io.FileIO(path, "wb")
downloader = MediaIoBaseDownload(fh, request)
done = False
while done is False:
_, done = downloader.next_chunk()
print("Downloaded %s" % metadata['name'])
do_upload(path, service)
# An alternative method would be to use partial download headers
# and convert and upload the parts individually. Perhaps a
# future release will implement this.
# Mode sets the mode of updating 0 > Verbose, 1 > Notification, 2 > silent
def update(self, mode=0, opts=None):
items = self.api.list_files(opts)
if not items:
print('No UDS files were found.')
elif mode != 2: # Duplicate silent...
table = []
with open("data.txt", 'w') as init: # Create data.txt if it does not exist
init.write("{")
init.write("}")
for item in items: # Show online data in form Encoded >> Size to show drive size and on disk size in order
record = [item.name, item.encoded_size, item.size, item.id_]
with open("data.txt", 'r') as data3: # Read data stored offline
user_data = json.load(data3)
temp_name = str(item.name)
user_data[temp_name] = item.id_
with open("data.txt", 'w') as data4: # Write data to data.txt
json.dump(user_data, data4, indent=3)
table.append(record)
with open("User.txt", 'w') as user:
user.write(tabulate(table, headers=[
'Name', 'Encoded', 'Size', 'ID']))
if mode == 0: # Verbose
print(tabulate(table, headers=[
'Name', 'Encoded', 'Size', 'ID']))
elif mode == 1: # Notify
print("Data Updated!\n")
def list(self, opts=None):
"""List UDS files
Prints a list of all UDS files. If a query is given, only the files
that match that query will be printed.
Args:
opts (str): Command line arguments
"""
items = self.api.list_files(opts)
if not items:
print('No UDS files were found.')
else:
# print('\nUDS Files in Drive:')
table = [[item.name, item.size, item.encoded_size, item.id_] for item in items]
print(tabulate(table, headers=[
'Name', 'Size', 'Encoded', 'ID', ]))
# Alpha command to erase file via name
def erase(self, name, default=1, mode_=None, fallback=None):
if fallback is not None:
self.delete_file(fallback, name=name, mode_=mode_)
else:
with open("data.txt", 'r') as list_:
data_pull = json.load(list_)
list_.close()
id_ = data_pull[name]
self.delete_file(id_, name=name, mode_=mode_)
self.update(mode=default) # Updates files in data after being altered
def grab(self, name, default=1, fallback=None): # Alpha command to pull files via name
self.update(mode=default) # Sets update mode
if fallback is not None:
self.build_file(parent_id=fallback)
print()
else:
with open("data.txt", 'r') as list_: # Load ID values based on file name
data_pull = json.load(list_)
parent_id = data_pull[name] # Loads ID based on name
self.build_file(parent_id)
print()
def batch(self, part, opts=None): # Alpha command to bulk download based on part of a file name
self.update(mode=1) # Sets update mode
items = self.api.list_files(opts)
name_space = [] # List of names based on user part
id_space = []
check = 0
for item in items: # Checks if part is in the name of any UDS file and adds them to queue
if str(part) != "?":
if str(part) in str(item.name): # The name check
name_space.append(item.name)
id_space.append(item.id_)
check += 1
elif str(part) == "?":
name_space.append(item.name)
id_space.append(item.id_)
for i in range(check):
self.grab(fallback=id_space[i], name=name_space[i], default=2)
# Downloads the bulk using data and names
for names in range(len(name_space)):
# Update data, not necessary
self.grab(name_space[names], default=2)
# Alpha command to bulk upload files based on file name part
def bunch(self, file_part, path='.'):
files = os.listdir(path) # Make list of all files in directory
files_upload = []
for name in files: # Cycles through all files
if file_part != "?":
if file_part in name: # Checks if part is in any files and adds to list
files_upload.append(name)
elif file_part == "?":
files_upload.append(name)
for name_data in range(len(files_upload)): # Upload all files put in list
full_path = str(path) + "/" + str(files_upload[name_data])
self.do_chunked_upload(full_path)
print()
self.update(mode=1) # Necessary update to data
def wipe(self, part, opts=None): # Alpha command to bulk delete files based on file name part
self.update(mode=2) # Sets update mode
items = self.api.list_files(opts)
name_space = []
id_space = []
check = 0
for item in items: # Add names to list
if str(part) != "?":
if str(part) in str(item.name): # add names if they have part in them
name_space.append(item.name)
check += 1
id_space.append(item.id_)
elif str(part) == "?":
name_space.append(item.name)
check += 1
id_space.append(item.id_)
else:
print()
for i in range(check):
self.erase(fallback=id_space[i], name=name_space[i], default=2)
def hash_file(self, path):
sha = hashlib.md5()
with open(path, 'rb') as f:
while True:
data = f.read(CHUNK_READ_LENGTH_BYTES)
if not data:
break
sha.update(data)
return sha.hexdigest()
def actions(self, action, args):
switcher = {
"list": self.list,
"push": self.do_chunked_upload,
"pull": self.build_file,
"delete": self.delete_file
}
switcher.get(action)(args)
def get_downloads_folder():
if not os.path.exists(DOWNLOADS_FOLDER):
os.makedirs(DOWNLOADS_FOLDER)
return DOWNLOADS_FOLDER
def characters_to_bytes(chars):
return round((3 / 4) * chars)
def write_status(status):
sys.stdout.write("\r%s" % status)
sys.stdout.flush()
def _parse_args(empty=False):
"""Parse command line arguments"""
formatter = lambda prog: argparse.HelpFormatter(prog, max_help_position=52)
parser = argparse.ArgumentParser(formatter_class=formatter)
parser.add_argument("--push", metavar='path_to_file', nargs=1,
help="Uploads a file from this computer")
parser.add_argument("--bunch", metavar=('word_in_file', 'path_to_file'),
nargs='+', help="Uploads files from this computer")
parser.add_argument("--pull", metavar='id', nargs=1,
help="Downloads a UDS file")
parser.add_argument("-g", "--grab", metavar='name', nargs=1
,help="Downloads a UDS file")
parser.add_argument("-b", "--batch", metavar='word_in_file', nargs=1,
help="Downloads UDS files")
parser.add_argument("-l", "--list", metavar='query', const='', nargs='?',
help="Finds all UDS files")
parser.add_argument("-u", "--update", action='store_true',
help="Update cached UDS data")
parser.add_argument("-d", "--delete", metavar='id', nargs=1,
help="Deletes a UDS file")
parser.add_argument("-e", "--erase", metavar='name', nargs=1,
help="Deletes a UDS file")
parser.add_argument("-w", "--wipe", metavar='word_in_file', nargs=1,
help="Deletes UDS files")
parser.add_argument("-c", "--convert", metavar='id', nargs=1,
help="Converts UDS files")
parser.add_argument("-C", "--clear", action='store_true',
help="Clear file after conversion")
parser.add_argument("-D", "--disable-multi", action='store_false',
help="Disable multithreading")
if empty:
parser.print_help()
return None
return parser.parse_args()
def main():
global BASE_FOLDER, USE_MULTITHREADED_UPLOADS, DELETE_FILE_AFTER_CONVERT
uds = UDS()
# Initial look for folder and first time setup if not
uds_root = uds.api.get_base_folder()
BASE_FOLDER = uds_root
# Options
if len(sys.argv) < 2:
_parse_args(True)
sys.exit(1)
args = _parse_args()
USE_MULTITHREADED_UPLOADS = args.disable_multi
if args.push:
uds.do_chunked_upload(args.push[0])
if args.bunch:
if len(args.bunch) > 1:
uds.bunch(args.bunch[0], args.bunch[1])
else:
uds.bunch(args.bunch[0])
if args.pull:
uds.build_file(args.pull[0])
if args.grab:
uds.grab(args.grab[0])
if args.batch:
uds.batch(args.batch[0])
if args.list != None:
# to differentiate from case of '' where no query is provided
uds.list(args.list)
if args.update:
uds.update()
if args.delete:
uds.delete_file(args.delete[0])
if args.erase:
uds.erase(args.erase[0])
if args.wipe:
uds.wipe(args.wipe[0])
if args.convert:
DELETE_FILE_AFTER_CONVERT = args.clear
convert_file(args.convert[0])
def ext_upload_chunked_part(chunk):
"""Upload a chunked part to drive and return the size of the chunk"""
_api = GoogleAPI()
# print("Chunk %s, bytes %s to %s" % (chunk.part, chunk.range_start, chunk.range_end))
with open(chunk.path) as fd:
mm = mmap.mmap(fd.fileno(), 0, access=mmap.ACCESS_READ)
chunk_bytes = mm[chunk.range_start:chunk.range_end]
encoded_chunk = encoder.encode(chunk_bytes)
file_metadata = {
'name': chunk.media.name + str(chunk.part),
'mimeType': 'application/vnd.google-apps.document',
'parents': [chunk.parent],
'properties': {
'part': str(chunk.part)
}
}
mediaio_file = MediaIoBaseUpload(io.StringIO(encoded_chunk),
mimetype='text/plain')
_api.upload_single_file(mediaio_file, file_metadata)
return len(chunk_bytes)
if __name__ == '__main__':
main()
