#! /usr/bin/python
# -*- coding: UTF-8 -*-
"""
1.    对称加密算法使用AES256。
2.    对称加密算法的工作模式，使用GCM模式。
3.    加密中所需的IV值，由安全随机数发生器产生。对于AES分组加密算法GCM模式，IV的长度为建议12字节。
4.    额外认证数据（AAD）是可选参数， GCM模式仅保护AAD完整性而不对其提供机密性保护，如果有AAD数据，则加密和解密时必须提供相同的AAD数据，产品在
对敏感数据加密的同时有其它需要保护完整性但无需保护机密性的数据时（例如程序上下文数据），可以将其作为AAD参数
5.    Tag是GCM模式中的消息认证码，用于验证被加密数据及AAD的完整性，GCM的Tag建议为16字节，至少12字节，应用需要更短的Tag时，需满足
NIST SP 800-38D附录D中关于密文及AAD的长度和以及同一密钥下加密算法调用次数的限制条件
"""

import sys
import platform
import secrets
import hashlib
from Crypto.Cipher import AES
from base64 import urlsafe_b64encode, urlsafe_b64decode
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Hash import (
                             SHA1, SHA256, SHA224, SHA384,
                             SHA512, HMAC
                             )

pf_os, pf_name, pf_version, _, _, _ = platform.uname()
sys.path.append("/usr/lib64/python2.7/site-packages" + pf_version.split(".")[-1])

HASH_ALGOS = {
    'sha1': SHA1,
    'sha224': SHA224,
    'sha256': SHA256,
    'sha384': SHA384,
    'sha512': SHA512
}
HASH_CNT = 1000  # Number of hashes to compute one SHA256 takes 15 microsec,
SK_HASH_CNT = 10000  # Number of hashes to compute source key
SK_DK_LENGTH = 16  # Number of SK DK length
SALT_LENGTH = 16  # Length for the Password salt for PBKDF
HASH_ALGO = 'sha256'  # For PBKDF HMAC
IV_LENGTH = 12  # Length of GCM IV
TAG_LENGTH = 16  # Length of the GCM tag, truncate if larger than this
HASH_FUNC = HASH_ALGOS[HASH_ALGO]


def hmac256(secret, m):
    return HMAC.new(key=secret, msg=m, digestmod=HASH_FUNC).digest()

#######################################################
# key:秘钥 16 (AES-128), 24 (AES-192), or 32 (AES-256)#
# plaintext: 需要加密的数据                           #
# associated_data: 额外认证数据（ADD），可选参数      #
#######################################################


def hash256(key):
    return hashlib.sha256(key)


def _encrypt(key, plaintext, associated_data=''):
    # Generate a random 96-bit IV.
    with open("/dev/random", 'rb') as file:
        iv = file.read(IV_LENGTH)
    # 16 (AES-128), 24 (AES-192), or 32 (AES-256)
    if len(key) not in (16, 24, 32):
        key = hash256(key)  # makes it 256-bit
    # Construct an AES-GCM Cipher object with the given key and a
    # randomly generated IV.
    # create the ciphertext
    encryptor = AES.new(key=key, mode=AES.MODE_GCM, nonce=iv)

    # associated_data will be authenticated but not encrypted,
    # it must also be passed in on decryption.
    encryptor.update(associated_data.encode())

    ctx = encryptor.encrypt(plaintext)
    # Encrypt the plaintext and get the associated ciphertext.
    # GCM does not require padding.
    tag = encryptor.digest()
    return (iv, ctx, tag)

##################################################
# key:16 (AES-128), 24 (AES-192), or 32 (AES-256)#
# ciphertext: 需要解密的数据                     #
# associated_data: 额外认证数据（ADD），可选参数 #
##################################################


def _decrypt(key, iv, ciphertext, tag, associated_data=''):
    # Construct a Cipher object, with the key, iv, and additionally the
    # GCM tag used for authenticating the message.
    if len(key) not in (16, 24, 32):
        key = hash256(key)  # makes it 256-bit

    decryptor = AES.new(key=key, mode=AES.MODE_GCM, nonce=iv)

    # We put associated_data back in or the tag will fail to verify
    # when we finalize the decryptor.
    decryptor.update(associated_data.encode())
    plaintext = decryptor.decrypt(ciphertext)

    # Decryption gets us the authenticated plaintext.
    # If the tag does not match an InvalidTag exception will be raised.
    decryptor.verify(tag)
    return plaintext


def pwencrypt(pw, m):
    """Encrypt the message m under pw using AES-GCM method (AEAD scheme).
    iv = 0   # Promise me you will never reuse the key
    c = <hash_style>.<iteration>.<urlsafe-base64 <salt><iv><tag><ctx>>
    :hash_style: sha-256 or sha-512, scrypt
    :iteration: Number of iteration. These two are the parameters
    for PBKDF2.
    Size of the ciphertext:
    """
    m = m.encode('utf-8', errors='ignore')
    itercnt = secrets.choice(list(range(HASH_CNT, 2*HASH_CNT)))
    header_txt = HASH_ALGO + '.' + str(itercnt)
    with open("/dev/random", 'rb') as file:
        sa = file.read(SALT_LENGTH)
    key = PBKDF2(
        pw, sa,
        dkLen=16,
        count=itercnt,
        prf=hmac256
    )
    iv, ctx, tag = _encrypt(key, m, associated_data=header_txt)
    # Salt (SALT_LENGTH), IV (IV_LENGTH), TAG (TAG_LENGTH)
    ctx_b64 = urlsafe_b64encode(sa + iv + tag + ctx)
    ctx_b64 = ctx_b64.decode() if isinstance(ctx_b64, bytes) else ctx_b64
    return header_txt + '.' + ctx_b64


def hmac_tmp(secret, m):
    return HMAC.new(key=secret, msg=m, digestmod=HASH_ALGOS[HASH_ALGO]).digest()


def gsecrtkey():
    """generate securite key
    """
    with open("/dev/random", 'rb') as file:
        sa = file.read(SALT_LENGTH)
    key = PBKDF2(
        sa, sa,
        dkLen=SK_DK_LENGTH,
        count=SK_HASH_CNT,
        prf=hmac_tmp
    )
    return urlsafe_b64encode(key)


def encrypt(k, m):
    """Symmetric encrpyt.
    """
    if not isinstance(m, bytes):
        m = bytes(m)
    iv, ctx, tag = _encrypt(k, m)
    return urlsafe_b64encode(iv + ctx + tag)


def decrypt(k, ctx):
    """Symmetric decrpyt.
    """
    if not isinstance(ctx, bytes):
        ctx = bytes(ctx)
    ctx_bin = urlsafe_b64decode(bytes(ctx))
    iv, c, tag = ctx_bin[:12], ctx_bin[12:-16], ctx_bin[-16:]
    try:
        return _decrypt(k, iv, c, tag)
    except Exception as e:
        print(e)
        raise(ValueError(e))


def sxor(s1, s2):
    # convert strings to a list of character pair tuples
    # go through each tuple, converting them to ASCII code (ord)
    # perform exclusive or on the ASCII code
    # then convert the result back to ASCII (chr)
    # merge the resulting array of characters as a string
    tmp_str = ''.join(chr(ord(a) ^ ord(b)) for a, b in zip(s1, s2))
    return urlsafe_b64encode(tmp_str.encode())
