# -*- coding: utf-8 -*-
    '''
        author:         Gupern
        createdDate:    2018-3-28
        lastModified:   2018-3-28 
    '''
    import argparse
    from pyDes import *
    import sys 
    from Cryptodome.Cipher import AES, ARC4
    from binascii import b2a_hex, a2b_hex, hexlify
    import sys,os,hashlib,time,base64
    import rc4
    import random, base64
    from hashlib import sha1
    from Crypto.Cipher import ARC4 as rc4cipher
    import rsa
    import sys,os,hashlib,time,base64

    class aes_prpcrypt():
        def __init__(self, key):

            # self.key = key
            self.key = bytes(key, encoding='utf-8')
            self.mode = AES.MODE_CBC

        #加密函数，如果text不是16的倍数【加密文本text必须为16的倍数！】，那就补足为16的倍数
        def encrypt(self, text):
            print(type(text))
            cryptor = AES.new(self.key, self.mode, self.key)
            print(type(text))

            #这里密钥key 长度必须为16（AES-128）、24（AES-192）、或32（AES-256）Bytes 长度.目前AES-128足够用
            length = 16
            count = len(text)
            add = length - (count % length)
            text = text + (bytes('\0',encoding="utf-8") * add)
            self.ciphertext = cryptor.encrypt(text)
            #因为AES加密时候得到的字符串不一定是ascii字符集的，输出到终端或者保存时候可能存在问题
            #所以这里统一把加密后的字符串转化为16进制字符串
            return b2a_hex(self.ciphertext)

        #解密后，去掉补足的空格用strip() 去掉
        def decrypt(self, text):
            cryptor = AES.new(self.key, self.mode, self.key)
            plain_text = cryptor.decrypt(a2b_hex(text))
            return plain_text.rstrip(bytes('\0', encoding="utf-8"))

    #rsa加密  
    def rsaEncrypt(data, pubkey):  

        #明文编码格式  
        content = data
        #公钥加密  
        crypto = rsa.encrypt(content,pubkey)  
        return crypto

    #rsa解密  
    def rsaDecrypt(data,pk):  
        #私钥解密  
        content = rsa.decrypt(data,pk)  
        con=content 
        return con  

    def des_algorithm(encrypt_or_decrypt, data, key1, key2):
        '''des de/encrypt algorithm'''
        k = des(key1, CBC, key2, pad=None, padmode=PAD_PKCS5)

        if encrypt_or_decrypt=="encrypt":
            res = k.encrypt(data)
            return True, res
        elif encrypt_or_decrypt=="decrypt":
            res = k.decrypt(data)
            return True, res
        else:
            print("methods is wrong")
            return False, ''

    def aes_algorithm(encrypt_or_decrypt, data, key1, key2):
        '''des de/encrypt algorithm'''
        k = aes_prpcrypt(key1)

        if encrypt_or_decrypt=="encrypt":
            res = k.encrypt(data)
            return True, res
        elif encrypt_or_decrypt=="decrypt":
            res = k.decrypt(data)
            return True, res
        else:
            print("methods is wrong")
            return False, ''

    def rc4_algorithm(encrypt_or_decrypt, data, key1, key2):
        '''des de/encrypt algorithm'''
        if encrypt_or_decrypt=="encrypt":
            key = bytes(key1, encoding='utf-8')
            enc = rc4cipher.new(key)
            res = enc.encrypt(data)
            return True, res
        elif encrypt_or_decrypt=="decrypt":
            key = bytes(key1, encoding='utf-8')
            enc = rc4cipher.new(key)
            res = enc.decrypt(data)
            return True, res
        else:
            print("methods is wrong")
            return False, ''

    def rsa_algorithm(encrypt_or_decrypt, data, key1, key2):
        '''rsa de/encrypt algorithm'''
        if encrypt_or_decrypt=="encrypt":
            #生成公钥、私钥  1024位
            (pubkey, privkey) = rsa.newkeys(1024)  
            # print("public key is: ", pubkey)
            # print("private key is: ", privkey)
            with open('./public.pem','w+') as f:
                f.write(pubkey.save_pkcs1().decode())
            with open('./private.pem','w+') as f:
                f.write(privkey.save_pkcs1().decode())

            res = b''
            dataLen = len(data) # 分块加密
            while dataLen>0:
                # print(dataLen)
                dataLen -= 117
                crypto = rsaEncrypt(data[:117], pubkey)
                data = data[117:]
                res += crypto
                print(len(res))
            return True, res
        elif encrypt_or_decrypt=="decrypt":
            dataLen = len(data)
            res = b''
            with open('./private.pem','r') as f:
                privkey = rsa.PrivateKey.load_pkcs1(f.read().encode())
            key2 = privkey
            # 分组解密，公钥为1024时，加密时117位一组，解密时128位一组
            while dataLen>0:
                print(dataLen)
                dataLen -= 128
                crypto = rsaDecrypt(data[:128], key2)
                data = data[128:]
                res += crypto
            return True, res
        else:
            print("methods is wrong")
            return False, ''


    if __name__ == '__main__':
        parser = argparse.ArgumentParser()
        parser.add_argument("input_file_location", help="Input the file location you want to process")
        parser.add_argument("encrypt_or_decrypt", help="you want to encrypt or decrypt?")
        parser.add_argument("key1", help="key1, for des")
        parser.add_argument("key2", help="key2, can be none")
        parser.add_argument("output_file_location", help="the location you want to output")
        parser.add_argument("methods", help="input the algorithm you wanna, des/aes/rc4/rsa")


        args = parser.parse_args()

        input_file_location = sys.argv[1]
        encrypt_or_decrypt = sys.argv[2]
        key1 = sys.argv[3]
        key2 = sys.argv[4]
        output_file_location = sys.argv[5]
        methods = sys.argv[6]
        data = open(input_file_location, 'rb').read()

        print(type(data))
        # print('data:', data, '\n')

        if methods=='des':
            success, res = des_algorithm(encrypt_or_decrypt, data, key1, key2)
        elif methods=='aes':
            success, res = aes_algorithm(encrypt_or_decrypt, data, key1, key2)
        elif methods=='rc4':
            success, res = rc4_algorithm(encrypt_or_decrypt, data, key1, key2)
        elif methods=='rsa':
            success, res = rsa_algorithm(encrypt_or_decrypt, data, key1, key2)


        # print('res:', res, '\n', type(res))
        # 加密时
        if encrypt_or_decrypt=='encrypt':
            with open(output_file_location, 'wb') as f:
                f.write(res)

        # 解密时
        elif encrypt_or_decrypt=='decrypt':
            with open(output_file_location, 'wb') as f:
                f.write(res)