gapis.py

#! /usr/bin/env python

import httplib2
from httplib import BadStatusLine
from dateutil import parser

from apiclient.discovery import build
from apiclient.http import MediaFileUpload
from apiclient import errors
from oauth2client.client import OAuth2WebServerFlow
from oauth2client.client import SignedJwtAssertionCredentials
from oauth2client.client import AccessTokenRefreshError

from commons import *

# single user
def create_drive_service_2_steps(client_id, client_secrect, oauth_scope, redirect_uri):
    # Run through the OAuth flow and retrieve credentials
    flow = OAuth2WebServerFlow(client_id, client_secrect, oauth_scope, redirect_uri)
    authorize_url = flow.step1_get_authorize_url()
    print 'Go to the following link in your browser: ' + authorize_url
    code = raw_input('Enter verification code: ').strip()
    credentials = flow.step2_exchange(code)

    # Create an httplib2.Http object and authorize it with our credentials
    http = httplib2.Http()
    http = credentials.authorize(http)

    try:
        drive_service = build('drive', 'v2', http=http)
        return drive_service
    except AccessTokenRefreshError, error:
        print 'Error when get drive service: %s' % error

    return None

# domain-wide
def create_drive_service(service_account_pkcs12_file,\
                        service_account_email, scope, user_email):
    """Build and returns a Drive service object authorized with the service accounts
        that act on behalf of the given user.

    Args:
        user_email: The email of the user.
    Returns:
        Drive service object.
    """
    f = file(service_account_pkcs12_file, 'rb')
    key = f.read()
    f.close()

    credentials = SignedJwtAssertionCredentials(service_account_email, key,\
                        scope=scope, sub=user_email)
    print "Finish getting credentials for user %s" % user_email

    http = httplib2.Http()
    http = credentials.authorize(http)

    print "Finish authorize user %s" % user_email

    try:
        drive_service = build('drive', 'v2', http=http)
        return drive_service
    except AccessTokenRefreshError, error:
        print "Error when getting drive service of user %s:\n > Error: %s"\
                        % (user_email, error)

    return None


######################### File utils #######################################

# get file by id
def get_file(service, file_id):
    """Print a file's metadata.

    Args:
        service: Drive API service instance.
        file_id: ID of the file to print metadata for.
    """
    try:
        file = service.files().get(fileId=file_id).execute()
        return file
    except BadStatusLine, badstatus:
        print 'Error when getting file by id: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'An error occurred: %s' % error

    return None


# Folder: mimeType = 'application/vnd.google-apps.folder'
# Text file: mimeType = 'text/plain'
def insert_file(service, title, description, mime_type, filename, parent_id=None):
    """Insert new file.

    Args:
        service: Drive API service instance.
        title: Title of the file to insert, including the extension.
        description: Description of the file to insert.
        parent_id: Parent folder's ID.
        mime_type: MIME type of the file to insert.
        filename: Filename of the file to insert.
    Returns:
        Inserted file metadata if successful, None otherwise.
    """
    media_body = MediaFileUpload(filename, mimetype=mime_type, resumable=True)
    body = {
        'title': title,
        'description': description,
        'mimeType': mime_type
    }

    # Set the parent folder.
    if parent_id:
        body['parents'] = [{'id': parent_id}]

    try:
        file = service.files().insert(
                    body=body,
                    media_body=media_body).execute()
        return file['id']
    except BadStatusLine, badstatus:
        print 'Error when inserting file: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'Insert file error: %s' % error

    return None


def retrieve_all(service):
    """Retrieve a list of File resources.

    Args:
    service: Drive API service instance.
    Returns:
    List of File resources.
    """

    result = []
    page_token = None
    while True:
        try:
            param = {}
            if page_token:
                param['pageToken'] = page_token
            files = service.files().list(**param).execute()

            result.extend(files['items'])
            page_token = files.get('nextPageToken')
            if not page_token:
                break
        except BadStatusLine, badstatus:
            print 'Error when retrieving all files: %s' % badstatus
            break
        except errors.HttpError, error:
            print 'An error occurred: %s' % error
            break

    return result


def retrieve_perms(service, obj_id):
    """Retrieve a list of permissions.
    Args:
        service: Drive API service instance.
        file_id: ID of the file to retrieve permissions for.
    Returns:
        List of permissions.
    """
    try:
        permissions = service.permissions().list(fileId=obj_id).execute()
        return permissions.get('items', [])
    except BadStatusLine, badstatus:
        print 'Error when retrieving permissions: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'Retrieve permissions error: %s' % error

    return None


def insert_perm(service, obj_id, value, perm_type, role, additionalRoles=[]):
    """Insert a new permission.

    Args:
        service:    Drive API service instance.
        file_id:    ID of the file to insert permission for.
        value:      User or group e-mail address, domain name or None for 'default'
                    type.
        perm_type:  The value 'user', 'group', 'domain' or 'default'.
        role:       The value 'owner', 'writer' or 'reader'.
    Returns:
        The inserted permission if successful, None otherwise.
    """

    new_permission = {
        'value': value,
        'type': perm_type,
        'role': role
    }
    if additionalRoles:
        new_permission['additionalRoles'] = additionalRoles

    try:
        perm = service.permissions().insert( \
                    fileId=obj_id, body=new_permission, sendNotificationEmails=False).execute()
        return perm['id']
    except BadStatusLine, badstatus:
        print 'Error when inserting permission: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'Insert permission error: %s' % error

    return None



# copy file permissions across domains using service account
def copy_perm(src_service, dest_service, src_user_email, dest_user_email, src_obj_id, dest_obj_id):
    src_perms = retrieve_perms(src_service, src_obj_id)
    if src_perms:
        for perm in src_perms:
            if 'emailAddress' in perm.keys():
                # if only the user is in our domain
                src_domain = src_user_email.split('@')[1]
                if src_domain in perm['emailAddress']:
                    dest_domain = dest_user_email.split('@')[1]
                    value = switch_email_domain(perm['emailAddress'], dest_domain)
                else:
                    value = perm['emailAddress']

                if perm['emailAddress'] != src_user_email:
                    if perm['emailAddress'] != dest_user_email:
                        perm_type = perm['type']
                        role = perm['role']
                        additionalRoles = []
                        if 'additionalRoles' in perm.keys():
                            additionalRoles = perm['additionalRoles']
                        print "Copy perm role %s for %s to file %s of user %s" \
                                % (role, perm['emailAddress'], dest_obj_id, dest_user_email)
                        insert_perm(dest_service, dest_obj_id, value, perm_type, role, additionalRoles)



def remove_perm(service, obj_id, permission_id):
    """Remove a permission.

    Args:
        service: Drive API service instance.
        file_id: ID of the file to remove the permission for.
        permission_id: ID of the permission to remove.
    """
    try:
        service.permissions().delete(
            fileId=obj_id, permissionId=permission_id).execute()
    except BadStatusLine, badstatus:
        print 'Error when removing permission: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'Remove permissions error: %s' % error



def get_perm_id_for_email(service, email):
    """Prints the Permission ID for an email address.

    Args:
        service: Drive API service instance.
        email: Email address to retrieve ID for.
    """
    try:
        id_resp = service.permissions().getIdForEmail(email=email).execute()
        # print id_resp['id']
        return id_resp['id']
    except BadStatusLine, badstatus:
        print 'Error when getting permission id for email: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'An error occured: %s' % error

    return None



def get_existed_files(service, org_obj_id, parentid='root'):
    same_title_files = []
    is_existed = False
    is_old = False

    org_file = get_file(service, org_obj_id)

    # If parentid is provided, search in that folder,
    # If no, search in root only by default
    query_string = "'%s' in parents \
                    and title = '%s' \
                    and trashed = false" \
                    % (parentid, org_file['title'])

    same_title_files = search_files(service, query_string)
    if same_title_files:
        is_existed = True
        for s_file in same_title_files:
            # if s_file['parents'][0]['id']==
            if s_file['mimeType'] == org_file['mimeType']:
                if s_file['modifiedDate'] < org_file['modifiedDate']:
                    is_old = True
                    break

    return same_title_files, is_existed, is_old


def delete_file(service, file_id):
    """Permanently delete a file, skipping the trash.

    Args:
        service: Drive API service instance.
        file_id: ID of the file to delete.
    """
    try:
        service.files().delete(fileId=file_id).execute()
    except BadStatusLine, badstatus:
        print 'Error when deleting file: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'Delete file error: %s' % error


def search_files(service, query_string):
    result = []
    page_token = None
    while True:
        try:
            param = {}
            if page_token:
                param['pageToken'] = page_token
            param['q'] = query_string
            files = service.files().list(**param).execute()

            result.extend(files['items'])
            page_token = files.get('nextPageToken')
            if not page_token:
                break
        except BadStatusLine, badstatus:
            print 'Error when searching files: %s' % badstatus
            break
        except errors.HttpError, error:
            print 'Error when searching file: %s' % error
            break

    return result


def retrieve_shared_with_me_files(service):
    query_string = 'sharedWithMe'
    return search_files(service, query_string)



def retrieve_own_files(service):
    tmp = []
    allfiles = search_files(service, "trashed = false")
    if allfiles:
        for file in allfiles:
            if file['owners'][0]['isAuthenticatedUser']:
                tmp.append(file)

    return tmp

def get_own_files_by_email(service, user_email):
    query = "'%s' in owners and trashed = false" % user_email
    files = search_files(service, query)
    return files


def get_own_files(service):
    query = "'me' in owners and trashed =false"
    files = search_files(service, query)
    print "\nGot %s files\n" % len(files)
    return files


def get_parents(service, file_id):
    """Print a file's parents.
    Args:
        service: Drive API service instance.
        file_id: ID of the file to print parents for.
    """
    try:
        parents = service.parents().list(fileId=file_id).execute()
        #~ for parent in parents['items']:
          #~ print 'File Id: %s' % parent['id']
        return parents
    except BadStatusLine, badstatus:
        print 'BadStatusLine error when getting parents of file (%s): %s' % (file_id, badstatus)
    except errors.HttpError, error:
        print 'HttpError when getting parents of file (%s): %s' % (file_id, error)

    return None


def rename_file(service, file_id, new_title):
    """Rename a file.

    Args:
    service: Drive API service instance.
    file_id: ID of the file to rename.
    new_title: New title for the file.
    Returns:
    Updated file metadata if successful, None otherwise.
    """
    try:
        file = {'title': new_title}

        # Rename the file.
        updated_file = service.files().patch(
                            fileId=file_id,
                            body=file,
                            fields='title').execute()

        return updated_file
    except BadStatusLine, badstatus:
        print 'Error when renaming file: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'Rename file error: %s' % error

    return None


def update_file(service, file_id, new_title, new_description, new_mime_type,
                    new_filename, new_revision):
    """Update an existing file's metadata and content.

    Args:
        service: Drive API service instance.
        file_id: ID of the file to update.
        new_title: New title for the file.
        new_description: New description for the file.
        new_mime_type: New MIME type for the file.
        new_filename: Filename of the new content to upload.
        new_revision: Whether or not to create a new revision for this file.
    Returns:
        Updated file metadata if successful, None otherwise.
    """
    try:
        # First retrieve the file from the API.
        file = service.files().get(fileId=file_id).execute()

        # File's new metadata.
        file['title'] = new_title
        file['description'] = new_description
        file['mimeType'] = new_mime_type

        # File's new content.
        media_body = MediaFileUpload(
            new_filename, mimetype=new_mime_type, resumable=True)

        # Send the request to the API.
        updated_file = service.files().update(
            fileId=file_id,
            body=file,
            newRevision=new_revision,
            media_body=media_body).execute()
        return updated_file
    except BadStatusLine, badstatus:
        print 'Error when updating file: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'An error occurred: %s' % error

    return None


# make a copy of a file on a same account
def copy_file(service, org_file, parentid=None):
    """Copy an existing file.

    Args:
        service: Drive API service instance.
        origin_file_id: ID of the origin file to copy.
        copy_title: Title of the copy.

    Returns:
        The copied file if successful, None otherwise.
    """
    body = {}
    body['title'] = org_file['title']
    body['mimeType'] = org_file['mimeType']

    # if the copying file is a child of a folder
    if parentid:
        body['parents'] = [{'id': parentid}]

    try:
        copied_file = service.files().copy(fileId=org_file['id'], body=body).execute()
        return copied_file
    except BadStatusLine, badstatus:
        print 'Error when copying file: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'Copy file error: %s' % error

    return None


def copy_unique_file(service, org_file, parentid=None):
    print "Copying file %s of parentid %s" % (org_file['id'], parentid)
    #~ org_title = org_file['title']
    org_title = clean_query_string(org_file['title'])
    query = "'me' in owners and title = '%s' and trashed = false" \
                % org_title

    if parentid:
        query += " and '%s' in parents" % parentid
    else:
        query += " and 'root' in parents"
    existed_files = search_files(service, query)
    if existed_files:
        tmp = {}
        for file in existed_files:
            if org_file['mimeType'] == file['mimeType'] \
                or (org_file['mimeType'] == 'application/vnd.google-apps.spreadsheet' \
                    and file['mimeType'] == 'application/vnd.google-apps.form'):

                if is_older(file, org_file):
                    print "Delete existed file %s" % (file['title'].encode('utf8'))
                    delete_file(service, file['id'])
                else:
                    #~ print "Skip copying file %s" % org_file['title'].encode('utf8')
                    tmp[parser.parse(file['modifiedDate'])] = file

        if len(tmp) >= 2: #rename duplicate files and return the unchanged file
            print "Skip copying file and rename duplicate files"
            return rename_dup_files_by_modified_date(service, tmp)
        elif len(tmp) == 1:
            print "Skip copying file"
            return tmp.values()[0]['id']

    copied_file = copy_file(service, org_file, parentid)

    print "Finish copying file %s" % (org_file['id'])

    if copied_file:
        return copied_file['id']

    return None


#########################################################################

###################### Folder ###########################################


def insert_folder(service, title, desc, parentid=None):
    body = {
        'title': title,
        'description': desc,
        'mimeType': 'application/vnd.google-apps.folder'
    }
    if parentid:
        body['parents'] = [{'id': parentid}]

    try:
        folder = service.files().insert(body=body).execute()
        # pprint.pprint(folder)
        return folder['id']
    except BadStatusLine, badstatus:
        print 'Error when inserting folder: %s' % badstatus
        # break
    except errors.HttpError, error:
        print 'Insert folder error: %s' % error
    return None


def copy_folder(service, folder_id, folder_title, parentid=None):
    # 1. create a folder with the same name in mydrive
    # 2. make a copy of all files in the source folder
    # 3. Assign the new folder as parents of the copied files

    new_created_ids = []
    new_folderid = insert_folder(service, folder_title, folder_title, parentid)
    new_created_ids.append({'src_id': folder_id, 'dest_id': new_folderid})

    query_string = "'%s' in parents" % (folder_id)

    files = search_files(service, query_string)
    for file in files:
        if file['mimeType'] == 'application/vnd.google-apps.folder':
            sub_created_ids = copy_folder(service, file['id'], file['title'], parentid=new_folderid)
            if sub_created_ids:
                new_created_ids += sub_created_ids
        else:
            copied_fileid = copy_unique_file(service, file, parentid=new_folderid)
            if copied_fileid:
                new_created_ids.append({'src_id': file['id'], 'dest_id': copied_fileid})

    return new_created_ids



def copy_unique_folder(service, folder_id, folder_title, parentid=None):
    # 1. create a folder with the same name in mydrive
    # 2. make a copy of all files in the source folder
    # 3. Assign the new folder as parents of the copied files

    new_created_ids = []

    folder_title = clean_query_string(folder_title)
    # copy the folder
    # check if there is any existed folder
    query = "'me' in owners and title = '%s' \
                and mimeType = 'application/vnd.google-apps.folder' \
                and trashed = false" \
                % (folder_title)
    if parentid:
        query += " and '%s' in parents" % parentid
    else:
        query += " and 'root' in parents"
    existed_folders = search_files(service, query)
    if existed_folders:
        new_folderid = existed_folders[0]['id']
    else:
        new_folderid = insert_folder(service, folder_title, folder_title, parentid)

    if new_folderid:
        new_created_ids.append({'src_id': folder_id, 'dest_id': new_folderid})

        # copy the children files and folders
        query_string = "'%s' in parents and trashed = false" \
                        % (folder_id)
        files = search_files(service, query_string)
        if files:
            print "All children of folder %s" % folder_title.encode('utf8')
            for file in files:
                print "+ %s" % file['title'].encode('utf8')
                if file['mimeType'] == 'application/vnd.google-apps.folder':
                    sub_created_ids = copy_unique_folder(service, file['id'], file['title'], parentid=new_folderid)
                    if sub_created_ids:
                        new_created_ids += sub_created_ids
                else:
                    copied_fileid = copy_unique_file(service, file, parentid=new_folderid)
                    if copied_fileid:
                        new_created_ids.append({'src_id': file['id'], 'dest_id': copied_fileid})
            print "Finish copying children of folder %s\n" % folder_title.encode('utf8')

    return new_created_ids


########################################################################

################## Rename duplicate files ##############################

# dup_files_dict = { datetime_obj1: file1, datetime_obj2: file2,}
def rename_dup_files_by_modified_date(service, dup_files_dict):
    order = dup_files_dict.keys()
    order.sort(reverse=True) # keys list in DESC order
    for i in order:
        if order.index(i) > 0:
            print "Renaming file %s" % dup_files_dict[i]['title']
            new_title = dup_files_dict[i]['title'] + " (" \
                                + str(order.index(i)) + ")"
            updated_file = rename_file(service, dup_files_dict[i]['id'], new_title)
            if updated_file:
                print "File %s has been renamed to %s" % (dup_files_dict[i]['title'], new_title)
            else:
                print "Fail to rename file %s" % dup_files_dict[i]['title']

    return dup_files_dict[order[0]]['id']


# rename duplicate files of a single user
def rename_all_dup_files(service):
    files = get_own_files(service)
    if files:
        filename_list = get_unique_file_name_list(files)
        if len(filename_list) < len(files):
            for fn in filename_list:
                dup_files_dict = {}
                for file in files:
                    if file['mimeType'] == fn['mimeType']:
                        if file['parents']:
                            if file['parents'][0]['id'] == fn['parentid']:
                                if file['title'] == fn['title']:
                                    dt = parser.parse(file['modifiedDate'])
                                    dup_files_dict[dt] = file
                if len(dup_files_dict) > 1:
                    rename_dup_files_by_modified_date(service, dup_files_dict)

########################################################################




# main operations for migration

# files_map = [ {'src': 'src_email@domain.com', 'dest': 'dest_email@domain.com', 'files': fileslist}, {}, ]
def share_files(service, files_map):
    perm_type = 'user'
    role = 'reader'

    perms_list = []
    shared_files = []
    for fm in files_map:
        for file in fm['files']:
            # only share files and folders in root
            if file['parents']:
                if file['parents'][0]['isRoot']:
                    permid = insert_perm(service, file['id'], fm['dest'], perm_type, role)
                    if permid:
                        perm_pair = {'fileid': file['id'], 'permid': permid}
                        perms_list.append(perm_pair)
                        shared_files.append(file)

    return perms_list, shared_files

# make a copy of shared files
def make_a_copy(service, shared_files):
    new_files_map = []
    for file in shared_files:
        if file['parents']:
            if file['parents'][0]['isRoot']:
                if file['mimeType'] == 'application/vnd.google-apps.folder':
                    # have to check more
                    new_folderids = copy_unique_folder(service, file['id'], file['title'])
                    if new_folderids:
                        new_files_map += new_folderids
                else:
                    new_fileid = copy_unique_file(service, file)
                    if new_fileid:
                        new_files_map.append({'src_id': file['id'], 'dest_id': new_fileid})

    return new_files_map

# perms_list = [{'fileid': <fileid>, 'permid': <permid>}, {...},]
def disable_sharing(service, perms_list):
    for perm in perms_list:
        remove_perm(service, perm['fileid'], perm['permid'])


def copy_perms(src_service, dest_service, src_user_email, dest_user_email, new_files_map):
    for file_map in new_files_map:
        if file_map['dest_id']:
            copy_perm(src_service, dest_service, src_user_email,\
                        dest_user_email, file_map['src_id'], file_map['dest_id'])
        else:
            print "Missing dest_id when copying permissions of %s" % src_user_email