class ServerSideEncryptionRule(object):

    def __init__(self, sse_algorithm=None, kms_master_keyid=None, kms_data_encryption=None):

        self.sse_algorithm = sse_algorithm
        self.kms_master_keyid = kms_master_keyid
        self.kms_data_encryption = kms_data_encryption class PartInfo(object):
    """表示分片信息的文件。

    该文件既用于 :func:`list_parts <oss2.Bucket.list_parts>` 的输出，也用于 :func:`complete_multipart_upload
    <oss2.Bucket.complete_multipart_upload>` 的输入。

    :param int part_number: 分片号
    :param str etag: 分片的ETag
    :param int size: 分片的大小。用在 `list_parts` 的结果里，也用与分片对象做crc combine得到整个对象crc64值
    :param int last_modified: 该分片最后修改的时间戳，类型为int。参考 :ref:`unix_time`
    :param int part_crc: 该分片的crc64值
    """
    def __init__(self, part_number, etag, size=None, last_modified=None, part_crc=None):
        self.part_number = part_number
        self.etag = etag
        self.size = size
        self.last_modified = last_modified
        self.part_crc = part_crc SERVER_SIDE_ENCRYPTION_AES256 = 'AES256' SERVER_SIDE_ENCRYPTION_KMS = 'KMS' SERVER_SIDE_ENCRYPTION_SM4 = 'SM4' KMS_DATA_ENCRYPTION_SM4 = 'SM4' OSS_SERVER_SIDE_ENCRYPTION = "x-oss-server-side-encryption" OSS_SERVER_SIDE_ENCRYPTION_KEY_ID = "x-oss-server-side-encryption-key-id" OSS_SERVER_SIDE_DATA_ENCRYPTION = "x-oss-server-side-data-encryption" def determine_part_size(total_size,
                        preferred_size=None):
    """确定分片上传是分片的大小。

    :param int total_size: 总共需要上传的长度
    :param int preferred_size: 用户期望的分片大小。如果不指定则采用defaults.part_size

    :return: 分片大小
    """
    if not preferred_size:
        preferred_size = defaults.part_size

    return _determine_part_size_internal(total_size, preferred_size, defaults.max_part_count) class SizedFileAdapter(object):
    """通过这个适配器（Adapter），可以把原先的 `file_object` 的长度限制到等于 `size`。"""
    def __init__(self, file_object, size):
        self.file_object = file_object
        self.size = size
        self.offset = 0

    def read(self, amt=None):
        if self.offset >= self.size:
            return ''

        if (amt is None or amt < 0) or (amt + self.offset >= self.size):
            data = self.file_object.read(self.size - self.offset)
            self.offset = self.size
            return data

        self.offset += amt
        return self.file_object.read(amt)

    @property
    def len(self):
        return self.size
from .common import *
from oss2.models import ServerSideEncryptionRule, PartInfo
from oss2 import (SERVER_SIDE_ENCRYPTION_KMS, SERVER_SIDE_ENCRYPTION_AES256,
                  SERVER_SIDE_ENCRYPTION_SM4, KMS_DATA_ENCRYPTION_SM4)
from oss2.headers import (OSS_SERVER_SIDE_ENCRYPTION, OSS_SERVER_SIDE_ENCRYPTION_KEY_ID,
                          OSS_SERVER_SIDE_DATA_ENCRYPTION)
from oss2 import determine_part_size, SizedFileAdapter
        result = bucket.init_multipart_upload(key, headers=headers)
        ret_sse_algo = result.headers.get(OSS_SERVER_SIDE_ENCRYPTION)
        ret_data_algo = result.headers.get(OSS_SERVER_SIDE_DATA_ENCRYPTION)
        ret_kms_key_id = result.headers.get(OSS_SERVER_SIDE_ENCRYPTION_KEY_ID)
        self.assertEqual(SERVER_SIDE_ENCRYPTION_KMS, ret_sse_algo)
        self.assertEqual(KMS_DATA_ENCRYPTION_SM4, ret_data_algo)
        self.assertIsNotNone(ret_kms_key_id)

        kms_key_id = ret_kms_key_id
        upload_id = result.upload_id
        parts = []

        # Upload part with the encyrption headers will be failed.
        headers = dict()
        headers[OSS_SERVER_SIDE_ENCRYPTION] = SERVER_SIDE_ENCRYPTION_KMS
        headers[OSS_SERVER_SIDE_DATA_ENCRYPTION] = KMS_DATA_ENCRYPTION_SM4
        with open(filename, 'rb') as fileobj:
            part_number = 1
            num_to_upload = part_size
            self.assertRaises(oss2.exceptions.InvalidArgument, bucket.upload_part, key, upload_id, part_number,
                                            SizedFileAdapter(fileobj, num_to_upload), headers=headers)

        # Upload part with none encryption headers.
        with open(filename, 'rb') as fileobj:
            part_number = 1
            offset = 0
            while offset < total_size:
                num_to_upload = min(part_size, total_size - offset)

                result = bucket.upload_part(key, upload_id, part_number,
                                            SizedFileAdapter(fileobj, num_to_upload))

                parts.append(PartInfo(part_number, result.etag))

                offset += num_to_upload
                part_number += 1

                ret_sse_algo = result.headers.get(OSS_SERVER_SIDE_ENCRYPTION)
                ret_data_algo = result.headers.get(OSS_SERVER_SIDE_DATA_ENCRYPTION)
                ret_kms_key_id = result.headers.get(OSS_SERVER_SIDE_ENCRYPTION_KEY_ID)
                self.assertEqual(SERVER_SIDE_ENCRYPTION_KMS, ret_sse_algo)
                self.assertEqual(KMS_DATA_ENCRYPTION_SM4, ret_data_algo)
                self.assertEqual(kms_key_id, ret_kms_key_id)

        # Complete multipart upload with encryption headers.
        self.assertRaises(oss2.exceptions.InvalidArgument, bucket.complete_multipart_upload, key, upload_id, parts, headers=headers)

        # Complete multipart upload with none encryption headers.
        result = bucket.complete_multipart_upload(key, upload_id, parts)
        self.assertEqual(result.status, 200)
        ret_sse_algo = result.headers.get(OSS_SERVER_SIDE_ENCRYPTION)
        ret_data_algo = result.headers.get(OSS_SERVER_SIDE_DATA_ENCRYPTION)
        ret_kms_key_id = result.headers.get(OSS_SERVER_SIDE_ENCRYPTION_KEY_ID)
        self.assertEqual(SERVER_SIDE_ENCRYPTION_KMS, ret_sse_algo)
        self.assertEqual(KMS_DATA_ENCRYPTION_SM4, ret_data_algo)
        self.assertEqual(kms_key_id, ret_kms_key_id)

        bucket.delete_object(key)


    def test_resumable_uoload(self):
        auth = oss2.Auth(OSS_ID, OSS_SECRET)
        bucket_name = OSS_BUCKET + "-test-resumable-upload-data-encryption"
        bucket = oss2.Bucket(auth, self.endpoint, bucket_name)
        bucket.create_bucket()

        small_object = 'requestpayment-test-resumable-upload-small-object'
        big_object = 'requestpayment-test-resumable-upload-big-object'

        # Create tmp file smaller than multipart_threshold
        file_name = self._prepare_temp_file_with_size(150 * 1024)

        # Resumale upload small object
        headers = dict()
        headers[OSS_SERVER_SIDE_ENCRYPTION] = SERVER_SIDE_ENCRYPTION_KMS
        headers[OSS_SERVER_SIDE_DATA_ENCRYPTION] = KMS_DATA_ENCRYPTION_SM4
        result = oss2.resumable_upload(bucket, small_object, file_name,
                                       multipart_threshold=(200 * 1024), num_threads=2, part_size=(100 * 1024),
                                       headers=headers)
        ret_sse_algo = result.headers.get(OSS_SERVER_SIDE_ENCRYPTION)
        ret_data_algo = result.headers.get(OSS_SERVER_SIDE_DATA_ENCRYPTION)
        ret_kms_key_id = result.headers.get(OSS_SERVER_SIDE_ENCRYPTION_KEY_ID)
        self.assertEqual(SERVER_SIDE_ENCRYPTION_KMS, ret_sse_algo)
        self.assertEqual(KMS_DATA_ENCRYPTION_SM4, ret_data_algo)
        self.assertIsNotNone(ret_kms_key_id)
        self.assertEqual(result.status, 200)
        bucket.delete_object(small_object)

        # Start big file test
        # Create big file bigger than multipart_threshold
        file_name = self._prepare_temp_file_with_size(11 * 1024 * 1024)

        headers = dict()
        headers[OSS_SERVER_SIDE_ENCRYPTION] = SERVER_SIDE_ENCRYPTION_KMS
        headers[OSS_SERVER_SIDE_DATA_ENCRYPTION] = KMS_DATA_ENCRYPTION_SM4
        result = oss2.resumable_upload(bucket, big_object, file_name,
                                       multipart_threshold=(200 * 1024), num_threads=2, part_size=(100 * 1024),
                                       headers=headers)

        self.assertEqual(result.status, 200)
        self.assertEqual(SERVER_SIDE_ENCRYPTION_KMS, ret_sse_algo)
        self.assertEqual(KMS_DATA_ENCRYPTION_SM4, ret_data_algo)
        self.assertIsNotNone(ret_kms_key_id)
        self.assertEqual(result.status, 200)
        bucket.delete_object(big_object)

    def test_copy_object(self):
        auth = oss2.Auth(OSS_ID, OSS_SECRET)
        bucket_name = OSS_BUCKET + "-test-copy-object-data-encryption"
        bucket = oss2.Bucket(auth, self.endpoint, bucket_name)
        bucket.create_bucket()

        object_name = 'test-copy-object-src'
        data = b'a' * 1024
        user_header_key = 'x-oss-meta-user1'
        user_header_value = 'user_value'

        headers = dict()
        headers[user_header_key] = user_header_value
        headers[OSS_SERVER_SIDE_ENCRYPTION] = SERVER_SIDE_ENCRYPTION_KMS

        headers[OSS_SERVER_SIDE_DATA_ENCRYPTION] = KMS_DATA_ENCRYPTION_SM4