gmail download and save on disk encrypted

gmail_downloader.py

#!/usr/bin/env python

"""
command line program to download gmail inbox messages between specified dates
"""

from __future__ import print_function
import argparse
import pprint
import base64
from datetime import datetime
import sys
import json
from Crypto.Cipher import AES
from Crypto import Random

from apiclient import discovery, errors
from httplib2 import Http
from oauth2client import file, client, tools

__author__ = 'xuefeng huang'

def valid_date(s):
    try:
        datetime.strptime(s, "%Y/%m/%d")
        return s
    except ValueError:
        msg = "Not a valid date: '{0}'.".format(s)
        raise argparse.ArgumentTypeError(msg)


parser = argparse.ArgumentParser(parents=[tools.argparser], description='program to save gmail messaage on disk encrypted')
parser.add_argument('-s', "--startdate", help="The start date - format YYYY/MM/DD ", 
                    required=True, type=valid_date)
parser.add_argument('-e', "--enddate", help="The end date - format YYYY/MM/DD ", 
                    required=True, type=valid_date)

flags = parser.parse_args()

if datetime.strptime(flags.startdate, "%Y/%m/%d") >= datetime.strptime(flags.enddate, "%Y/%m/%d"):
    sys.exit('start date must be earlier than end date')

def get_message(service, msg_id):
    try:
        if msg_id is not None:
            message = service.users().messages().get(userId='me', id=msg_id, format='raw').execute()
            msg_str = base64.urlsafe_b64decode(message['raw'].encode('ASCII'))
            return msg_str

    except errors.HttpError as error:
        print('An error occurred: {}'.format(error)) 



def get_all_messages(service, query):
    try:
        # see https://developers.google.com/gmail/api/v1/reference/users/messages/get for gmail python API
        response = service.users().messages().list(userId='me',
                                                   q=query).execute()
        messages = []
        if 'messages' in response:
          messages.extend(response['messages'])

        while 'nextPageToken' in response:
          page_token = response['nextPageToken']
          response = service.users().messages().list(userId='me', q=query,
                                             pageToken=page_token).execute()
          messages.extend(response['messages'])

        message_list = []
        for item in messages:
            msg = get_message(service, item.get('id', None))
            print('saving message:\n{}'.format(msg))
            print('=' * 20)
            message_list.append({'message': msg})

        return message_list 


    except errors.HttpError as error:
        print('An error occurred: {}'.format(error)) 

def encrypt_msg(msg):
    key, iv= generate_key()
    # msg padding to 16 bytes
    msg = padding(msg)
    cipher = AES.new(key, AES.MODE_CBC, iv)

    # save encrypted file
    with open('msg_encrypted.txt', 'w') as f:
        f.write(base64.b64encode(cipher.encrypt(msg)))

def generate_key():
    key = Random.new().read(16)
    iv = Random.new().read(AES.block_size)

    # save key and iv for decryption
    with open('key.txt', 'w') as f:
        f.write(base64.b64encode(key))

    with open('iv.txt', 'w') as f:
        f.write(base64.b64encode(iv))

    return key, iv

def padding(msg):
    length = 16 - (len(msg) % 16)
    msg += chr(length) * length # remove padding msg = msg[:-ord(msg[-1])]
    return msg

if __name__ == '__main__':
    SCOPES = 'https://www.googleapis.com/auth/gmail.readonly'
    store = file.Storage('storage.json')
    creds = store.get()
    if not creds or creds.invalid:
        flow = client.flow_from_clientsecrets('client_secret.json', SCOPES)
        # see https://developers.google.com/api-client-library/python/guide/aaa_oauth#commandline for doing this step
        creds = tools.run_flow(flow, store, flags)

    GMAIL = discovery.build('gmail', 'v1', http=creds.authorize(Http()))

    query = 'before: {0} after: {1}'.format(flags.enddate, flags.startdate)
    message_list = get_all_messages(GMAIL, query)
    msg_json_str = json.dumps(message_list)
    encrypt_msg(msg_json_str)