[Python] Cryptography 라이브러리 (hazmat 모듈) 사용법

안녕하세요, 오늘은 Cryptography 라이브러리에서 hazmat 모듈을 사용하는 방법에 대해서 알아보겠습니다.

 

hazmat 모듈을 사용하는 예시는 다음과 같습니다.

 

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import os
import base64

def encrypt_file(file_path, password, output_path):
    backend = default_backend()
    salt = os.urandom(16)
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=backend
    )
    key = kdf.derive(password.encode())

    iv = os.urandom(16)
    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=backend)
    encryptor = cipher.encryptor()

    with open(file_path, 'rb') as f:
        data = f.read()

    padder = padding.PKCS7(128).padder()
    padded_data = padder.update(data) + padder.finalize()

    encrypted_data = encryptor.update(padded_data) + encryptor.finalize()

    with open(output_path, 'wb') as f:
        f.write(base64.b64encode(salt + iv + encrypted_data))

# 사용 예시
encrypt_file('model.h5', 'test_password', 'model.encrypted')

 

위의 코드는 tensorflow2 프레임워크 내 특정 모델 model.h5 파일을 특정 비밀번호 = 'test_password'를 활용해서 model.encrypted 파일로 암호화 하는 코드 입니다.

---

이어서, 해당 파일을 복호화하는 코드에 대해서 알아보겠습니다.

 

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64

def decrypt_file(encrypted_file_path, password, output_path):
    backend = default_backend()

    with open(encrypted_file_path, 'rb') as f:
        encrypted_data = base64.b64decode(f.read())

    salt = encrypted_data[:16]
    iv = encrypted_data[16:32]
    cipher_text = encrypted_data[32:]

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=backend
    )
    key = kdf.derive(password.encode())

    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=backend)
    decryptor = cipher.decryptor()

    padded_data = decryptor.update(cipher_text) + decryptor.finalize()

    unpadder = padding.PKCS7(128).unpadder()
    data = unpadder.update(padded_data) + unpadder.finalize()

    with open(output_path, 'wb') as f:
        f.write(data)

# 사용 예시
decrypt_file('model.encrypted', 'test_password', 'model.h5')

 

---

추가적으로, 복호화 시 따로 model.h5 파일로 저장하지 않고 복호화 후, 메모리 상헤서 바로 로드하여 사용하는 방법에 대해서 알아보겠습니다.

 

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64
import io
import h5py
import tensorflow as tf

def decrypt_and_load_model(encrypted_file_path, password):
    backend = default_backend()

    with open(encrypted_file_path, 'rb') as f:
        encrypted_data = base64.b64decode(f.read())

    salt = encrypted_data[:16]
    iv = encrypted_data[16:32]
    cipher_text = encrypted_data[32:]

    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=backend
    )
    key = kdf.derive(password.encode())

    cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=backend)
    decryptor = cipher.decryptor()

    padded_data = decryptor.update(cipher_text) + decryptor.finalize()

    unpadder = padding.PKCS7(128).unpadder()
    data = unpadder.update(padded_data) + unpadder.finalize()

    # 메모리 상에서 파일 객체로 변환
    model_data = io.BytesIO(data)

    # TensorFlow 모델 로드
    with h5py.File(model_data, 'r') as f:
    	model = tf.keras.models.load_model(f)

    return model

# 사용 예시
model = decrypt_and_load_model('model.encrypted', 'test_password')

 

위의 코드에서 파일을 저장하지 않고, io.BytesIO를 사용하여 데이터를 메모리 상의 파일 객체로 만들고, tf.keras.models.load_model() 함수를 사용하여 Tensorflow 모델을 로드할 수 있습니다.

 

---

추가적으로 각 OS 별로 환경 변수 설정하는 방법에 대해서 알아보겠습니다.

 

1. bash를 사용하는 경우  : nano ~/.bashrc
2. zsh를 사용하는 경우 : nano ~/.zshrc

파일 끝에 다음과 같은 줄을 추가 합니다.

 

export ENCRYPTION_PASSWORD=test_password

 

이어서 변경 사항을 적용합니다.

 

source ~/.bashrc  # 또는 ~/.zshrc 등

 

환경 변수로 정상적으로 적용되었는 지 확인하기 위해서는 다음 명령어로 확인 가능합니다.

<Linux / MacOs>

echo $ENCRYPTION_PASSWORD

 

<Windows>

echo %ENCRYPTION_PASSWORD%

 

파이썬 코드에서 위와 같이 설정한 환경 변수를 사용하기 위해서는 다음과 같은 코드를 사용 가능합니다.

 

import os

# 환경 변수에서 비밀번호 가져오기
password = os.getenv('ENCRYPTION_PASSWORD')

if password is None:
    raise ValueError("환경 변수 ENCRYPTION_PASSWORD가 설정되지 않았습니다.")

# 사용 예시
model = decrypt_and_load_model('encrypted_model.encrypted', password)

---

지금까지, Cryptography 라이브러리 내 hazmat 모듈을 사용하여 파일을 암호화, 복호화 및 메모리 상의 파일 객체로 만드는 방법, 그리고 환경 변수 설정 및 사용방법에 대해서 알아보았습니다.

감사드립니다.