430 lines
16 KiB
Python
430 lines
16 KiB
Python
# -*- coding: utf-8 -*-
|
||
|
||
import os
|
||
import random
|
||
import logging
|
||
import copy
|
||
import base64
|
||
import struct
|
||
from .cos_comm import *
|
||
from Crypto.Cipher import AES
|
||
from Crypto.PublicKey import RSA
|
||
from Crypto import Random
|
||
from Crypto.Util import Counter
|
||
from Crypto.Cipher import PKCS1_OAEP, PKCS1_v1_5
|
||
from .cos_exception import CosClientError
|
||
from .streambody import StreamBody
|
||
from abc import ABCMeta, abstractmethod
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
_AES_CTR_COUNTER_BITS_LENGTH = 8 * 16
|
||
_AES_256_KEY_SIZE = 32
|
||
|
||
|
||
def random_key(key_len):
|
||
return Random.new().read(key_len)
|
||
|
||
|
||
def random_iv():
|
||
iv = Random.new().read(16)
|
||
return iv
|
||
|
||
|
||
def iv_to_big_int(iv):
|
||
iv_pair = struct.unpack(">QQ", iv)
|
||
iv_int = iv_pair[0] << 64 | iv_pair[1]
|
||
return iv_int
|
||
|
||
|
||
class AESCTRCipher(object):
|
||
"""数据加密类，用于加密用户数据"""
|
||
|
||
def __init__(self):
|
||
"""初始化"""
|
||
self.__block_size_len = AES.block_size
|
||
self.__cipher = None
|
||
self.__key_len = _AES_256_KEY_SIZE
|
||
|
||
def new_cipher(self, key, start, offset=0):
|
||
"""初始化AES加解密对象
|
||
|
||
:param key(string): 对称密钥
|
||
:param start(int): 对称加密初始随机值
|
||
:param offset(int): 主要用于解密，想要解密文本的偏移，必须为16的整数倍
|
||
"""
|
||
block_index_offset = self.__calc_offset(offset)
|
||
new_start = start + block_index_offset
|
||
my_counter = Counter.new(_AES_CTR_COUNTER_BITS_LENGTH, initial_value=new_start)
|
||
self.__cipher = AES.new(key, AES.MODE_CTR, counter=my_counter)
|
||
|
||
def encrypt(self, plaintext):
|
||
"""加密数据
|
||
|
||
:param plaintext(string): 需要加密的数据
|
||
"""
|
||
if self.__cipher is None:
|
||
raise CosClientError('cipher is not initialized')
|
||
return self.__cipher.encrypt(plaintext)
|
||
|
||
def decrypt(self, plaintext):
|
||
"""解密数据
|
||
|
||
:param plaintext(string): 需要解密的数据，数据的起始位置必须为16的整数倍
|
||
"""
|
||
if self.__cipher is None:
|
||
raise CosClientError('cipher is not initialized')
|
||
return self.__cipher.decrypt(plaintext)
|
||
|
||
# offset 必须为block_size的整数倍
|
||
def __calc_offset(self, offset):
|
||
"""通过文本偏移计算counter的偏移
|
||
|
||
:param offset(int): 文本偏移
|
||
"""
|
||
if not self.__is_block_aligned(offset):
|
||
raise CosClientError('offset is not align to encrypt block')
|
||
return offset // self.__block_size_len
|
||
|
||
def __is_block_aligned(self, offset):
|
||
"""判断文本偏移是否是block_size对齐
|
||
|
||
:param offset(int): 文本偏移
|
||
"""
|
||
if offset is None:
|
||
offset = 0
|
||
return 0 == offset % self.__block_size_len
|
||
|
||
def adjust_read_offset(self, offset):
|
||
"""用于调整读取的offset为block_size对齐"""
|
||
|
||
if offset:
|
||
offset = (offset // self.__block_size_len) * self.__block_size_len
|
||
return offset
|
||
|
||
def get_key(self):
|
||
"""获取随机密钥"""
|
||
return random_key(self.__key_len)
|
||
|
||
def get_counter_iv(self):
|
||
"""获取对称加密初始随机值"""
|
||
return random_iv()
|
||
|
||
|
||
class RSAKeyPair:
|
||
"""封装有公钥和私钥"""
|
||
|
||
def __init__(self, public_key, private_key):
|
||
self.publick_key = public_key
|
||
self.private_key = private_key
|
||
|
||
|
||
class RSAKeyPairPath:
|
||
"""封装有公钥和私钥的路径"""
|
||
|
||
def __init__(self, public_key_path, private_key_path):
|
||
self.public_key_path = public_key_path
|
||
self.private_key_path = private_key_path
|
||
|
||
|
||
class BaseProvider(object):
|
||
"""客户端主密钥加密基类"""
|
||
|
||
def __init__(self, cipher):
|
||
"""初始化
|
||
:param cipher(an AES object): 数据加解密类
|
||
"""
|
||
self.data_cipher = cipher
|
||
|
||
def get_data_key(self):
|
||
"""随机获取数据加解密密钥"""
|
||
return self.data_cipher.get_key()
|
||
|
||
@abstractmethod
|
||
def init_data_cipher(self):
|
||
"""初始化cipher"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
def init_data_cipter_by_user(self, encrypt_key, encrypt_iv, offset=0):
|
||
"""根据密钥初始化cipher"""
|
||
pass
|
||
|
||
def adjust_read_offset(self, start):
|
||
"""用于调整读取的offset为block_size对齐"""
|
||
return self.data_cipher.adjust_read_offset(start)
|
||
|
||
def make_data_encrypt_adapter(self, stream):
|
||
"""创建数据流加密适配器"""
|
||
size = 0
|
||
if hasattr(stream, '__len__'):
|
||
size = len(stream)
|
||
elif hasattr(stream, 'tell') and hasattr(stream, 'seek'):
|
||
current = stream.tell()
|
||
stream.seek(0, os.SEEK_END)
|
||
size = stream.tell()
|
||
stream.seek(current, os.SEEK_SET)
|
||
else:
|
||
return None
|
||
return DataEncryptAdapter(stream, size, copy.copy(self.data_cipher))
|
||
|
||
def make_data_decrypt_adapter(self, rt, offset):
|
||
"""创建数据流解密适配器"""
|
||
return DataDecryptAdapter(rt, copy.copy(self.data_cipher), offset)
|
||
|
||
|
||
class RSAProvider(BaseProvider):
|
||
"""客户端非对称主密钥加密类"""
|
||
|
||
def __init__(self, key_pair_info=None, cipher=AESCTRCipher(), passphrase=None):
|
||
"""初始化"""
|
||
super(RSAProvider, self).__init__(cipher=cipher)
|
||
|
||
default_rsa_dir = os.path.expanduser('~/.cos_local_rsa')
|
||
default_public_key_path = os.path.join(default_rsa_dir, '.public_key.pem')
|
||
default_private_key_path = os.path.join(default_rsa_dir, '.private_key.pem')
|
||
self.__encrypt_obj = None
|
||
self.__decrypt_obj = None
|
||
self.__data_key = None
|
||
self.__data_iv = None
|
||
|
||
if isinstance(key_pair_info, RSAKeyPair):
|
||
self.__encrypt_obj = PKCS1_v1_5.new(RSA.importKey(key_pair_info.publick_key, passphrase=passphrase))
|
||
self.__decrypt_obj = PKCS1_v1_5.new(RSA.importKey(key_pair_info.private_key, passphrase=passphrase))
|
||
elif isinstance(key_pair_info, RSAKeyPairPath):
|
||
self.__encrypt_obj, self.__decrypt_obj = self.__get_key_by_path(key_pair_info.public_key_path,
|
||
key_pair_info.private_key_path, passphrase)
|
||
else:
|
||
logger.info('key_pair_info is None, try to get key from default path')
|
||
if os.path.exists(default_private_key_path) and os.path.exists(default_public_key_path):
|
||
self.__encrypt_obj, self.__decrypt_obj = self.__get_key_by_path(default_public_key_path,
|
||
default_private_key_path, passphrase)
|
||
|
||
# 为用户自动创建rsa
|
||
if self.__encrypt_obj is None and self.__decrypt_obj is None:
|
||
logger.warning('fail to get rsa key, will generate key')
|
||
private_key = RSA.generate(2048)
|
||
public_key = private_key.publickey()
|
||
|
||
self.__encrypt_obj = PKCS1_OAEP.new(public_key)
|
||
self.__decrypt_obj = PKCS1_OAEP.new(private_key)
|
||
|
||
if not os.path.exists(default_rsa_dir):
|
||
os.makedirs(default_rsa_dir)
|
||
|
||
with open(default_private_key_path, 'wb') as f:
|
||
f.write(private_key.exportKey(passphrase=passphrase))
|
||
|
||
with open(default_public_key_path, 'wb') as f:
|
||
f.write(public_key.exportKey(passphrase=passphrase))
|
||
|
||
@staticmethod
|
||
def get_rsa_key_pair(public_key, private_key):
|
||
"""开放给用户，用于生成RSAKeyPair"""
|
||
if public_key is None or private_key is None:
|
||
raise CosClientError('public_key or private_key is not allowed to be None !!!')
|
||
return RSAKeyPair(public_key, private_key)
|
||
|
||
@staticmethod
|
||
def get_rsa_key_pair_path(public_key_path, private_key_path):
|
||
"""开放给用户，用于生成RSAKeyPairPath"""
|
||
if public_key_path is None or private_key_path is None:
|
||
raise CosClientError('public_key or private_key is not allowed to be None !!!')
|
||
return RSAKeyPairPath(public_key_path, private_key_path)
|
||
|
||
def __get_key_by_path(self, public_path=None, private_path=None, passphrase=None):
|
||
"""用于从提供的路径中获取公钥和私钥"""
|
||
if public_path is None or private_path is None:
|
||
return None, None
|
||
|
||
encrypt_obj, decrypt_obj = None, None
|
||
|
||
if os.path.exists(public_path) and os.path.exists(private_path):
|
||
with open(public_path, 'rb') as f:
|
||
encrypt_obj = PKCS1_OAEP.new(RSA.importKey(f.read(), passphrase=passphrase))
|
||
|
||
with open(private_path, 'rb') as f:
|
||
decrypt_obj = PKCS1_OAEP.new(RSA.importKey(f.read(), passphrase=passphrase))
|
||
|
||
return encrypt_obj, decrypt_obj
|
||
|
||
def init_data_cipher(self):
|
||
"""初始化cipher"""
|
||
encrypt_key = None
|
||
encrypt_iv = None
|
||
self.__data_key = self.get_data_key()
|
||
self.__data_iv = self.data_cipher.get_counter_iv()
|
||
start = iv_to_big_int(self.__data_iv)
|
||
self.data_cipher.new_cipher(self.__data_key, start)
|
||
encrypt_key = self.__encrypt_obj.encrypt(self.__data_key)
|
||
encrypt_iv = self.__encrypt_obj.encrypt(self.__data_iv)
|
||
return encrypt_key, encrypt_iv
|
||
|
||
def init_data_cipter_by_user(self, encrypt_key, encrypt_iv, offset=0):
|
||
"""根据密钥初始化cipher"""
|
||
self.__data_key = self.__decrypt_obj.decrypt(encrypt_key)
|
||
self.__data_iv = self.__decrypt_obj.decrypt(encrypt_iv)
|
||
start = iv_to_big_int(self.__data_iv)
|
||
self.data_cipher.new_cipher(self.__data_key, start, offset)
|
||
|
||
|
||
class AESProvider(BaseProvider):
|
||
"""客户端对称主密钥加密类"""
|
||
|
||
def __init__(self, aes_key=None, aes_key_path=None, cipher=AESCTRCipher()):
|
||
"""初始化"""
|
||
super(AESProvider, self).__init__(cipher=cipher)
|
||
self.__ed_obj = None
|
||
self.__data_key = None
|
||
self.__data_iv = None
|
||
self.__my_counter = Counter.new(_AES_CTR_COUNTER_BITS_LENGTH, initial_value=0)
|
||
self.__aes_key = aes_key
|
||
self.__aes_key_path = aes_key_path
|
||
|
||
def init_ed_obj(self):
|
||
default_aes_dir = os.path.expanduser('~/.cos_local_aes')
|
||
default_key_path = os.path.join(default_aes_dir, '.aes_key.pem')
|
||
if self.__aes_key:
|
||
aes_key = to_bytes(base64.b64decode(to_bytes(self.__aes_key)))
|
||
self.__ed_obj = AES.new(aes_key, AES.MODE_CTR, counter=self.__my_counter)
|
||
elif self.__aes_key_path:
|
||
if os.path.exists(self.__aes_key_path):
|
||
with open(self.__aes_key_path, 'rb') as f:
|
||
aes_key = f.read()
|
||
aes_key = to_bytes(base64.b64decode(to_bytes(aes_key)))
|
||
self.__ed_obj = AES.new(aes_key, AES.MODE_CTR, counter=self.__my_counter)
|
||
else:
|
||
logger.info('aes_key and aes_key_path is None, try to get key from default path')
|
||
if os.path.exists(default_key_path):
|
||
with open(default_key_path, 'rb') as f:
|
||
aes_key = f.read()
|
||
aes_key = to_bytes(base64.b64decode(to_bytes(aes_key)))
|
||
self.__ed_obj = AES.new(aes_key, AES.MODE_CTR, counter=self.__my_counter)
|
||
|
||
if self.__ed_obj is None:
|
||
logger.warning('fail to get aes key, will generate key')
|
||
aes_key = random_key(_AES_256_KEY_SIZE)
|
||
self.__ed_obj = AES.new(aes_key, AES.MODE_CTR, counter=self.__my_counter)
|
||
if not os.path.exists(default_aes_dir):
|
||
os.makedirs(default_aes_dir)
|
||
|
||
with open(default_key_path, 'wb') as f:
|
||
aes_key = to_bytes(base64.b64encode(to_bytes(aes_key)))
|
||
f.write(aes_key)
|
||
|
||
def init_data_cipher(self):
|
||
"""初始化cipher"""
|
||
encrypt_key = None
|
||
encrypt_iv = None
|
||
self.__data_key = self.get_data_key()
|
||
self.__data_iv = self.data_cipher.get_counter_iv()
|
||
start = iv_to_big_int(self.__data_iv)
|
||
self.data_cipher.new_cipher(self.__data_key, start)
|
||
|
||
self.init_ed_obj()
|
||
encrypt_key = self.__ed_obj.encrypt(self.__data_key)
|
||
encrypt_iv = self.__ed_obj.encrypt(self.__data_iv)
|
||
return encrypt_key, encrypt_iv
|
||
|
||
def init_data_cipter_by_user(self, encrypt_key, encrypt_iv, offset=0):
|
||
"""根据密钥初始化cipher"""
|
||
self.init_ed_obj()
|
||
self.__data_key = self.__ed_obj.decrypt(encrypt_key)
|
||
self.__data_iv = self.__ed_obj.decrypt(encrypt_iv)
|
||
start = iv_to_big_int(self.__data_iv)
|
||
self.data_cipher.new_cipher(self.__data_key, start, offset)
|
||
|
||
|
||
class MetaHandle(object):
|
||
"""用于获取/生成加密的元信息"""
|
||
|
||
def __init__(self, encrypt_key=None, encrypt_iv=None):
|
||
"""初始化
|
||
|
||
:param encrypt_key(string): 加密的数据密钥
|
||
:param encrypt_iv(bytes): 加密counter的初始值
|
||
"""
|
||
self.__encrypt_key = encrypt_key
|
||
self.__encrypt_iv = encrypt_iv
|
||
|
||
def set_object_meta(self, headers):
|
||
"""设置加密元信息到object的头部"""
|
||
meta_data = dict()
|
||
meta_data['x-cos-meta-client-side-encryption-key'] = to_bytes(base64.b64encode(to_bytes(self.__encrypt_key)))
|
||
meta_data['x-cos-meta-client-side-encryption-iv'] = to_bytes(base64.b64encode(to_bytes(self.__encrypt_iv)))
|
||
headers['Metadata'] = meta_data
|
||
return headers
|
||
|
||
def get_object_meta(self, headers):
|
||
"""从object的头部获取加密元信息"""
|
||
if 'x-cos-meta-client-side-encryption-key' in headers and 'x-cos-meta-client-side-encryption-iv' in headers:
|
||
self.__encrypt_key = base64.b64decode(to_bytes(headers['x-cos-meta-client-side-encryption-key']))
|
||
self.__encrypt_iv = base64.b64decode(to_bytes(headers['x-cos-meta-client-side-encryption-iv']))
|
||
return self.__encrypt_key, self.__encrypt_iv
|
||
|
||
|
||
class DataEncryptAdapter(object):
|
||
"""用于读取经过加密后的的数据"""
|
||
|
||
def __init__(self, data, content_len, data_cipher):
|
||
self._data = to_bytes(data)
|
||
self._data_cipher = data_cipher
|
||
self._content_len = content_len
|
||
self._read_len = 0
|
||
|
||
@property
|
||
def len(self):
|
||
return self._content_len
|
||
|
||
def read(self, length):
|
||
"""读取加密后的数据"""
|
||
if self._read_len >= self._content_len:
|
||
return ''
|
||
|
||
if length is None or length < 0:
|
||
bytes_to_read = self._content_len - self._read_len
|
||
else:
|
||
bytes_to_read = min(length, self._content_len - self._read_len)
|
||
|
||
if isinstance(self._data, bytes):
|
||
content = self._data[self._read_len:self._read_len + bytes_to_read]
|
||
else:
|
||
content = self._data.read(bytes_to_read)
|
||
|
||
self._read_len += bytes_to_read
|
||
content = self._data_cipher.encrypt(content)
|
||
return content
|
||
|
||
|
||
class DataDecryptAdapter(StreamBody):
|
||
"""用于读取经过解密后的数据"""
|
||
|
||
def __init__(self, rt, data_cipher, offset=0):
|
||
"""初始化
|
||
:param rt(request object): request请求返回的对象
|
||
:param data_cipher(an AES object): 数据加解密类
|
||
:param offset(int): 数据读取的偏移量
|
||
"""
|
||
super(DataDecryptAdapter, self).__init__(rt)
|
||
self._data_cipher = data_cipher
|
||
self._offset = offset
|
||
|
||
def read(self, length, auto_decompress=False):
|
||
"""读取解密后的数据"""
|
||
if self._read_len >= self._content_len:
|
||
return ''
|
||
|
||
if self._use_encoding and not auto_decompress:
|
||
content = self._rt.raw.read(length)
|
||
else:
|
||
try:
|
||
content = next(self._rt.iter_content(length))
|
||
except StopIteration:
|
||
return ''
|
||
|
||
content = self._data_cipher.decrypt(content)
|
||
if self._read_len < self._offset and self._read_len + len(content) > self._offset:
|
||
content = content[self._offset:]
|
||
self._read_len = self._offset
|
||
return content
|