# Path: testbench/csek.py
# def extract(request, is_source, context):
# def check(algorithm, key_b64, expected_sha256, context):
# def validation(request, expected_sha256, is_source, context):
# 
# Path: testbench/error.py
# class RestException(Exception):
#     def __init__(self, msg, code):
#     def as_response(self):
#     def handler(ex):
# def _simple_json_error(msg):
# def generic(msg, rest_code, grpc_code, context):
# def csek(context, rest_code=400, grpc_code=grpc.StatusCode.INVALID_ARGUMENT):
# def invalid(msg, context, rest_code=400, grpc_code=grpc.StatusCode.INVALID_ARGUMENT):
# def missing(name, context, rest_code=400, grpc_code=grpc.StatusCode.INVALID_ARGUMENT):
# def mismatch(
#     msg,
#     expect,
#     actual,
#     context,
#     rest_code=412,
#     grpc_code=grpc.StatusCode.FAILED_PRECONDITION,
# ):
# def notchanged(msg, context, rest_code=304, grpc_code=grpc.StatusCode.ABORTED):
# def notfound(name, context, rest_code=404, grpc_code=grpc.StatusCode.NOT_FOUND):
# def notallowed(context=None, rest_code=405, grpc_code=None):
# def already_exists(context=None):


import base64
import hashlib
import unittest
import grpc
from unittest.mock import ANY, Mock
from werkzeug.test import create_environ
from werkzeug.wrappers import Request
from google.storage.v2 import storage_pb2
from testbench import csek, error
    def test_extract_grpc_source(self):
        key_bytes = b"\001\002\003\004\005\006\007\010"
        key_hash = hashlib.sha256(key_bytes).digest()
        context = unittest.mock.Mock()
        algorithm, key, sha256 = csek.extract(
            storage_pb2.RewriteObjectRequest(
                copy_source_encryption_algorithm="AES256",
                copy_source_encryption_key_bytes=key_bytes,
                copy_source_encryption_key_sha256_bytes=key_hash,
            ),
            is_source=True,
            context=context,
        )
        self.assertEqual(algorithm, "AES256")
        self.assertEqual(key, base64.b64encode(key_bytes).decode("utf-8"))
        self.assertEqual(sha256, base64.b64encode(key_hash).decode("utf-8"))

    def test_check_success(self):
        key = b"1234567890" + b"1234567890" + b"1234567890" + b"AA"
        key_sha256 = hashlib.sha256(key).digest()
        key_b64 = base64.b64encode(key)
        expected = hashlib.sha256(key).digest()
        actual = csek.check("AES256", key_b64, key_sha256, None)
        self.assertEqual(expected, actual)

    def test_check_invalid_algorithm(self):
        key = b"1234567890" + b"1234567890" + b"1234567890" + b"AA"
        key_sha256 = hashlib.sha256(key).digest()
        key_b64 = base64.b64encode(key)
        with self.assertRaises(error.RestException) as rest:

# Path: testbench/csek.py
# def extract(request, is_source, context):
# def check(algorithm, key_b64, expected_sha256, context):
# def validation(request, expected_sha256, is_source, context):
# 
# Path: testbench/error.py
# class RestException(Exception):
#     def __init__(self, msg, code):
#     def as_response(self):
#     def handler(ex):
# def _simple_json_error(msg):
# def generic(msg, rest_code, grpc_code, context):
# def csek(context, rest_code=400, grpc_code=grpc.StatusCode.INVALID_ARGUMENT):
# def invalid(msg, context, rest_code=400, grpc_code=grpc.StatusCode.INVALID_ARGUMENT):
# def missing(name, context, rest_code=400, grpc_code=grpc.StatusCode.INVALID_ARGUMENT):
# def mismatch(
#     msg,
#     expect,
#     actual,
#     context,
#     rest_code=412,
#     grpc_code=grpc.StatusCode.FAILED_PRECONDITION,
# ):
# def notchanged(msg, context, rest_code=304, grpc_code=grpc.StatusCode.ABORTED):
# def notfound(name, context, rest_code=404, grpc_code=grpc.StatusCode.NOT_FOUND):
# def notallowed(context=None, rest_code=405, grpc_code=None):
# def already_exists(context=None):


# Path: tests/test_csek.py
import base64
import hashlib
import unittest
import grpc
from unittest.mock import ANY, Mock
from werkzeug.test import create_environ
from werkzeug.wrappers import Request
from google.storage.v2 import storage_pb2
from testbench import csek, error

    def test_extract_grpc_source(self):
        key_bytes = b"\001\002\003\004\005\006\007\010"
        key_hash = hashlib.sha256(key_bytes).digest()
        context = unittest.mock.Mock()
        algorithm, key, sha256 = csek.extract(
            storage_pb2.RewriteObjectRequest(
                copy_source_encryption_algorithm="AES256",
                copy_source_encryption_key_bytes=key_bytes,
                copy_source_encryption_key_sha256_bytes=key_hash,
            ),
            is_source=True,
            context=context,
        )
        self.assertEqual(algorithm, "AES256")
        self.assertEqual(key, base64.b64encode(key_bytes).decode("utf-8"))
        self.assertEqual(sha256, base64.b64encode(key_hash).decode("utf-8"))

    def test_check_success(self):
        key = b"1234567890" + b"1234567890" + b"1234567890" + b"AA"
        key_sha256 = hashlib.sha256(key).digest()
        key_b64 = base64.b64encode(key)
        expected = hashlib.sha256(key).digest()
        actual = csek.check("AES256", key_b64, key_sha256, None)
        self.assertEqual(expected, actual)

    def test_check_invalid_algorithm(self):
        key = b"1234567890" + b"1234567890" + b"1234567890" + b"AA"
        key_sha256 = hashlib.sha256(key).digest()
        key_b64 = base64.b64encode(key)
        with self.assertRaises(error.RestException) as rest:

            _ = csek.check("#### INVALID ####", key_b64, key_sha256, None)