mirror of
https://github.com/aykhans/AzSuicideDataVisualization.git
synced 2025-04-22 10:28:02 +00:00
61 lines
2.5 KiB
Python
61 lines
2.5 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
import base64
|
|
|
|
import pyarrow.parquet.encryption as pe
|
|
|
|
|
|
class InMemoryKmsClient(pe.KmsClient):
|
|
"""This is a mock class implementation of KmsClient, built for testing only.
|
|
"""
|
|
|
|
def __init__(self, config):
|
|
"""Create an InMemoryKmsClient instance."""
|
|
pe.KmsClient.__init__(self)
|
|
self.master_keys_map = config.custom_kms_conf
|
|
|
|
def wrap_key(self, key_bytes, master_key_identifier):
|
|
"""Not a secure cipher - the wrapped key
|
|
is just the master key concatenated with key bytes"""
|
|
master_key_bytes = self.master_keys_map[master_key_identifier].encode(
|
|
'utf-8')
|
|
wrapped_key = b"".join([master_key_bytes, key_bytes])
|
|
result = base64.b64encode(wrapped_key)
|
|
return result
|
|
|
|
def unwrap_key(self, wrapped_key, master_key_identifier):
|
|
"""Not a secure cipher - just extract the key from
|
|
the wrapped key"""
|
|
expected_master_key = self.master_keys_map[master_key_identifier]
|
|
decoded_wrapped_key = base64.b64decode(wrapped_key)
|
|
master_key_bytes = decoded_wrapped_key[:16]
|
|
decrypted_key = decoded_wrapped_key[16:]
|
|
if (expected_master_key == master_key_bytes.decode('utf-8')):
|
|
return decrypted_key
|
|
raise ValueError("Incorrect master key used",
|
|
master_key_bytes, decrypted_key)
|
|
|
|
|
|
def verify_file_encrypted(path):
|
|
"""Verify that the file is encrypted by looking at its first 4 bytes.
|
|
If it's the magic string PARE
|
|
then this is a parquet with encrypted footer."""
|
|
with open(path, "rb") as file:
|
|
magic_str = file.read(4)
|
|
# Verify magic string for parquet with encrypted footer is PARE
|
|
assert magic_str == b'PARE'
|