import os
import sys
import binascii

# Need to do some work with the path in order to get all the modules we need
# Get the current path of this file
curPath = os.path.dirname(os.path.realpath(__file__))
lambdaPath = os.path.join(curPath,"../awsPictureLambda/src")
sys.path.append(os.path.abspath(lambdaPath))
sys.path.append(os.path.abspath(os.path.join(curPath,"..")))
import lambda_function
from s3Backup import *

def isIvUnique(ivBytes):
    ivStr = binascii.hexlify(ivBytes).decode('utf-8')
    result = lambda_function.Picture.select().where(lambda_function.Picture.iv\
                                                    == ivStr)
    return len(result) == 0

def reEncrypt(oldKey, newKey, oldAuth, newAuth, pic):
    """Takes the information as specified in a picture description file and
    re-encrypts it"""
    outputFile = pic.file_name
    auth_tag = binascii.unhexlify(pic.auth_tag)
    success = True
    try:
        download_picture.getFile(pic.sha256, outputFile, bucket, oldKey, oldAuth, oldIv, auth_tag)
    except Exception as e:
        print('***ERROR: could not decrypt SHA: + ' + pic.sha256)
        success = False

    if success:
        # Now re-encrypt it
        encFile = outputFile + '.enc'
        cipherText = open(encFile, mode='wb')
        with open(outputFile, mode='rb') as plainText:
            (newIv, new_auth_tag) = file_crypto.encryptFile(newKey, newAuth,
                                                            plainText, cipherText,
                                                            ivCheckFunc=isIvUnique)
            print('IV: ' + binascii.hexlify(newIv).decode('utf-8'))
            print('Tag: ' + binascii.hexlify(new_auth_tag).decode('utf-8'))
        cipherText.close()

        with open(encFile, mode='rb') as encPic:
            ## Upload the new file to S3
            bucket.upload_fileobj(encPic, pic.sha256)

        ## Save the new IV and auth tag to the DB
        pic.iv = binascii.hexlify(newIv).decode('utf-8')
        pic.auth_tag = binascii.hexlify(new_auth_tag).decode('utf-8')
        pic.save()

    ### Clean up temp files
    if os.path.exists(pic.file_name):
        os.unlink(pic.file_name)
    if os.path.exists(encFile):
        os.unlink(encFile)

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='''Will re-encrypt all files,
                                     ensure the database is updated, and update
                                     S3. It is really mean to just change the
                                     IV that all files where encrypted with.
                                     You see I screwed up first and used an IV
                                     that was 256 bits. For GCM, this is fine,
                                     since the just compute the actual IV by
                                     runnign GHASH over the input. Once I tried
                                     to decrypt it with CTR things got weird.
                                     So, instead we should use a 96 bit IV as
                                     it turns out that using an IV like this is
                                     straight forward with GCM. Probably due to
                                     some oddity with multiplication in the
                                     Galois field.''')
    parser.add_argument('config',help='''Path to our aws JSON config file''')
    args = parser.parse_args()


    if 'DB_USER' not in os.environ or\
       'DB_PASS' not in os.environ or\
       'DB_ENDPT' not in os.environ:
           print('''ERROR: must specify DB_USER, DB_PASS, and DB_ENDPT
                 environment variable, just like the awsLambda function
                 expects''')
           sys.exit(1)

    # Now, get all the picture from the database
    pics = lambda_function.Picture.select()

    config = common_config.AWSConfig(args.config)
    (key, auth) = file_crypto.createKeyAndAuthData(config.get_key_phrase())
    (newKey, newAuth) = \
            file_crypto.createKeyAndAuthData_New(config.get_key_string())

    bucket = config.get_picture_bucket()
    maxP = len(pics)
    #TODO - need to re-run on element 18159 and 18209. These were skipped
    # earlier due to an off by one error. Duh
    #cur = 18160 # Originally failed here
    cur = 18210
    for pic in pics[cur:]:
        oldIv = binascii.unhexlify(pic.iv)
        reEncrypt(key, newKey, auth, newAuth, pic)
        cur += 1
        print(str(cur) + ' / ' + str(maxP))
