Python script to sync local folder to google drive

gdrive.py

#!/usr/bin/env python3

from __future__ import print_function
import os
import sys
import time
import json
import random
import inspect
import argparse
import hashlib
import datetime

# Google Drive
import pickle
from googleapiclient.discovery import build
from google_auth_oauthlib.flow import InstalledAppFlow
from google.auth.transport.requests import Request
from apiclient.http import MediaFileUpload


#--------------------------------------------------------------------
# Example use
#
# ./gdrive.py verify -c ~/credentials-gdrive-user.json -l ./test -r /test -RU
#

#--------------------------------------------------------------------
# Dependencies
#
# python3 -m pip install --upgrade google-api-python-client google-auth-httplib2 google-auth-oauthlib
#--------------------------------------------------------------------

g_silent = False

def Log(s):
    if g_silent:
        return
    fi = inspect.getframeinfo(inspect.stack()[1][0])
    print(fi.filename + "(" + str(fi.lineno) + "): " + str(s))
    sys.stdout.flush()

def clipStr(s, max, clipstart=False, clip='...'):
    if len(s) <= max:
        return s
    if clipstart:
        return clip + s[len(s) - max + len(clip):]
    return s[0:max-len(clip)] + clip

def fixStr(s, max, clipstart=False, clip='...', pad=' '):
    return clipStr(s, max, clipstart, clip).ljust(max, pad)

def showStatus(s, max=100):
    print("\r" + fixStr(s, max), end='')
    sys.stdout.flush()

# [ 87%] = {0:3.0f}
# [ 87.34%] = {0:6.2f}
def showProgress(p, s, max=100, clipstart=True, format="{0:3.0f}"):
    print("\r[" + format.format(p) + "%] " + fixStr(s, max, clipstart), end='')
    sys.stdout.flush()

def tsCheck(_p):
    now = datetime.datetime.now()
    if 'lastmin' not in _p or _p['lastmin'] != now.minute:
        _p['lastmin'] = now.minute
        print("\n--------- %s ---------\n" % (now.strftime('%b %d %H:%M')))


#--------------------------------------------------------------------
# Google Drive


def getCreds(cf, auth):

    creds = None
    scopes = [
                'https://www.googleapis.com/auth/drive.file',
                'https://www.googleapis.com/auth/drive.metadata.readonly'
             ]
    cache = './gdrive.token.cache'

    # Do we have saved credentials?
    if not auth and os.path.exists(cache):
        with open(cache, 'rb') as fh:
            creds = pickle.load(fh)

    # Login user if no credentials
    if not creds or not creds.valid:

        # Just need a refresh?
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())

        # Full request
        else:
            flow = InstalledAppFlow.from_client_secrets_file(cf, scopes)
            creds = flow.run_local_server(port=0)

        # Save the credentials for the next time
        with open(cache, 'wb') as fh:
            pickle.dump(creds, fh)

    return creds


def getGDrive(cf, auth):

    # Login
    creds = getCreds(cf, auth)
    if not creds:
        Log("Failed to load credentials")
        return None

    # Create service
    return build('drive', 'v3', credentials=creds)


def getItemInfo(gd, folder):

    # Folder components
    fp = folder.split('/')

    try:

        fi = gd.files().get(fileId='root').execute()
        if not fi:
            return None

        while True:

            cur = ''
            while not cur:
                if 0 >= len(fp):
                    return fi
                cur = fp.pop(0)

            res = gd.files().list(
                        q="'" + fi['id'] + "' in parents and name contains '" + cur.replace("'", r"\'") + "'",
                        spaces='drive',
                        pageSize=1,
                        fields="nextPageToken, files(id, name, size, parents)"
                        ).execute()
            fi = res.get('files', [])[0]

    except Exception as e:
        Log(e)

    return None


def listFiles(gd, folder='', id=''):

    files = []

    try:
        if not id:
            ii = getItemInfo(gd, folder)
            id = ii['id']

        max = 100
        page_token = None

        while True:

            # Pressure valve
            max -= 1
            if 0 >= max:
                break

            # Get a list of files
            res = gd.files().list(
                        q="'" + id + "' in parents",
                        pageToken=page_token,
                        spaces='drive',
                        pageSize=1000,
                        #fields="*"
                        fields="nextPageToken, files(id, name, size, parents, createdTime)"
                        ).execute()

            # Add files
            files += res.get('files', [])

            # Next page
            page_token = res.get('nextPageToken', None)
            if not page_token:
                break

    except Exception as e:
        Log(e)

    return files


def createFolder(gd, root, name):

    try:

        # Get root
        ri = getItemInfo(gd, root)
        if not ri or 'id' not in ri:
            Log("Invalid root : " + root)

        # Create the sub folder
        body = {
            'name': name,
            'parents': [ri['id']],
            'mimeType': 'application/vnd.google-apps.folder'
        }
        folder = gd.files().create(body=body, fields='id').execute()

        return folder

    except Exception as e:
        Log(e)

    return None

def uploadFile(gd, src, dst='', rid=''):

    if not os.path.isfile(src):
        Log("Invalid file : " + src)
        return False

    # Grab rid if needed
    if not rid:
        ri = getItemInfo(gd, dst)
        if not ri or 'id' not in ri:
            Log("Invalid remote root")
            return False
        rid = ri['id']

    fpath, fname = os.path.split(src)

    CHUNK_SIZE = 256 * 1024
    DEFAULT_CHUNK_SIZE = 100 * 1024 * 1024

    try:

        progress = 0
        media = MediaFileUpload(src, resumable=True, chunksize=CHUNK_SIZE)
        if not media:
            Log("Failed to create upload media object")
            return False

        res = gd.files().create( body = {'name': fname, 'parents': [rid]},
                                    media_body = media,
                                    fields = "id"
                                )

        if not res:
            Log("Invalid upload response object")
            return False


        response = None
        while response is None:

            # Process next file chunk
            status, response = res.next_chunk()

            if status:
                progress = float(100 * status.resumable_progress) / float(status.total_size)
                showProgress(progress, "Uploading  : %s" % clipStr(src, 85, True))

    except Exception as e:
        showProgress(progress, '!!! FAILED : ' + clipStr(src, 85, True))
        print('')
        return False

    showProgress(100, "Uploaded   : %s" % clipStr(src, 85, True))
    print('')

    return True


def calc_md5(fname):
    hash_md5 = hashlib.md5()
    with open(fname, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


def verifyFolder(_p, gd, loc, rem):

    tsCheck(_p)

    print("\n---------------------------------------------------------------------")
    print("--- %s -> %s\n" %(loc, rem))

    # Get root folder info
    ri = getItemInfo(gd, rem)
    if not ri or 'id' not in ri:
        Log("Invalid remote folder : " + rem)
        return False

    # Get remote files
    rfiles = listFiles(gd, id=ri['id'])
    if not rfiles:
        rfiles = []

    rmap = {}
    for v in rfiles:
        if 'name' in v:
            rmap[v['name']] = v

    # Get local files
    lfiles = os.listdir(loc)
    for v in lfiles:

        tsCheck(_p)

        # Skip hidden files
        if v[0] == '.':
            continue

        # Build full path
        fp = os.path.join(loc, v)

        # Is it a file?
        if os.path.isfile(fp):

            # Local file size
            lsz = int(os.path.getsize(fp))

            # Ignore zero size files
            if 0 >= lsz:
                Log("!!! File has zero size : " + fp)

            # Is it missing from the remote server?
            elif v not in rmap:
                if _p['upload']:

                    retry = 3
                    while 0 < retry:
                        retry -= 1
                        if not uploadFile(gd, fp, rid=ri['id']):
                            Log("!!! FAILED TO UPLOAD : " + fp)
                        else:
                            retry = -1

                else:
                    print("MISSING : " + fp)

            else:
                rsz = int(rmap[v]['size'])

                # Do the sizes match?
                if rsz != lsz:
                    print("SIZE ERROR : " + fp + " : " + str(rsz) + " != " + str(lsz))

                else:
                    pass
                    #Log("OK : " + fp)


        # Is it a directory?
        elif os.path.isdir(fp):

            # Is it missing?
            if v not in rmap:

                # Upload files
                if _p['upload']:
                    print("CREATE FOLDER : " + fp)
                    fi = createFolder(gd, rem, v)
                    if fi:
                        rmap[v] = fi

                else:
                    print("MISSING FOLDER : " + fp)

            # It's there
            if v in rmap:

                # Do we want to recurse into sub folders?
                if _p['recursive']:
                    verifyFolder(_p, gd, fp, os.path.join(rem, v))


#-------------------------------------------------------------------

def cmd_verify(_p):

    # Verify folder contents
    verifyFolder(_p, _p['gdrive'], _p['local'], _p['remote'])

def cmd_list(_p):

    gd = _p['gdrive']
    rem = _p['remote']

    ri = getItemInfo(gd, rem)
    if not ri or 'id' not in ri:
        Log("Invalid remote folder : " + _p['remote'])
        return False

    # Get remote files
    rfiles = listFiles(gd, id=ri['id'])
    if not rfiles:
        rfiles = []

    for v in rfiles:
        print(v['name'])


def main(_p):

    # Did we get any commands?
    if 0 >= len(_p['cmd']):
        Log("No commands given")
        return

    # Get google drive
    _p['gdrive'] = getGDrive(_p['creds'], _p['auth'])
    if not _p['gdrive']:
        Log("Failed to get GDrive object")
        return

    if _p['cmd'][0] == 'verify':
        return cmd_verify(_p)

    if _p['cmd'][0] == 'list':
        return cmd_list(_p)


if __name__ == '__main__':

    _p = vars()
    try:
        # Log("Los Geht's...")

        # Get command line params
        ap = argparse.ArgumentParser(description='Google Drive Sync')
        ap.add_argument('cmd', metavar='N', type=str, nargs='+', help='Commands')
        ap.add_argument('--creds', '-c', required=True, type=str, help='Credentials File')
        ap.add_argument('--remote', '-r', type=str, help='Remote path')
        ap.add_argument('--local', '-l', type=str, help='Local path')
        ap.add_argument('--recursive', '-R', action='store_true', help='Recursive')
        ap.add_argument('--upload', '-U', action='store_true', help='Upload new files')
        ap.add_argument('--download', '-D', action='store_true', help='Download new files')
        ap.add_argument('--auth', '-a', action='store_true', help='Force authentication')
        ap.add_argument('--silent', '-s', action='store_true', help='Turn off messages')
        _p = vars(ap.parse_args())

        # Silent mode
        if _p['silent']:
            g_silent = True

        # Show parameters
        Log("\n ---------------- Parameters -----------------\n"
            + json.dumps(_p, sort_keys=True, indent=2) + "\n")

        # Run
        _p['run'] = True
        main(_p)

    except KeyboardInterrupt:
        Log(" ~ keyboard ~ ")

    # except Exception as e:
    #     Log(" ~ exception ~ ")
    #     Log(str(e))

    finally:
        _p['run'] = False
        Log("Bye...")