Zenn
🤖

pickleファイルの秘匿化方法を考えてみた

2024/03/13に公開
Python
base64
pickle
暗号化
fernet
tech

はじめに

個人で作った機械学習モデルを誰かに渡す際、作った資産を流用されない様にするために、どの様な方法で防止することができそうか考えてみた。
思いの他、やりたい事を実現するのに時間がかかってしまったが、一応は最低限のガードはできたので、そのサンプルコードとハマった点をメモとして残しておく。

実行環境

【PC環境】
　　Windows 10 Pro　
【仮想環境のpython.Ver　および主要パッケージ情報】
　　Python 3.9.13
　　xgboost 2.0.3
　　scikit-learn 1.4.1.post1
　　cryptography 42.0.5

今回やりたい事のイメージ

適当な機械学習モデル(model.pkl)を作る。
これを使いスコアリングやクラス分類はできるが、モデルの中身は覗けない様にした形で提供したい。

以下の流れでサンプルコードを記載します。

  1. モデル構築とpickle保存
  2. モデルの保護
  3. 保護したモデルの使い方イメージ

1. モデル構築とpickle保存

まず、テスト用の資産（機械学習済みモデル）を準備

import pandas as pd
import xgboost as xgb
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# データを読み込む
iris = load_iris()
df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
df['class'] = iris.target
display(df)

X = df.drop(columns=['class'])
y = df['class']

# データをトレーニングセットとテストセットに分割
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# モデルのパラメータを設定
params = {
    'objective': 'multi:softmax',
    'num_class': 3,
    'eval_metric': 'merror'
}

# XGBoostのモデル構築
dtrain = xgb.DMatrix(X_train, label=y_train)
num_rounds = 100
model = xgb.train(params, dtrain, num_rounds)

# モデルの保存
pd.to_pickle(model, 'model.pkl')

もちろん、このままの状態では model.pkl を持ち出して使うこともできるし、pd.read_pickleなどで使用している特徴量も確認できてしまう。

2. モデルの保護

pickleファイル　→　base64でテキスト化　→　テキストを暗号化

# バイナリ形式のデータを読み込む
with open("model.pkl", "rb") as file:
    read_model = file.read()
print('<バイナリの状態>\r\n', read_model)

import base64

# PickleデータをBase64でエンコードする
encoded_data = base64.b64encode(read_model)
print('<base64でエンコードした状態>\r\n', encoded_data)

from cryptography.fernet import Fernet

# 鍵を作成して、base64でエンコードしたものを暗号化して txt形式で出力
passphrase = b'\xc8P\xde@\xf8\x0e\xfe\x81\xf5\xdf\x1b\xafLDe\xdf\xdfT\xd5\xab6\xa1\x9f\xf0\xd49\xb8\xce\x92A\x1c\x85'
key = base64.urlsafe_b64encode(passphrase)

f = Fernet(key)
encrypted_data = f.encrypt(encoded_data)
print('<暗号化した状態>\r\n', encrypted_data)

# 暗号化したモデル情報を txt形式で出力
with open("encrypted_model.txt", "wb") as file:
    file.write(encrypted_data)

上記コードのパスフレーズの設定と、keyの作成で超ハマりました。
っというのも、ChatGPTやググったりして出てくるサンプルコードは、key=Fernet.generate_key() の書き方しかなく、一旦ランダムで作った鍵をローカルに保持する形になってしまう。
そうなると、結局クライアントPC上にそのまま使える鍵(decode可能な鍵)が存在することになるので、鍵の在りかさえ分かれば復号ができてしまう。

鍵自体はローカルに置かず、モデル利用のタイミングで同じ鍵を作成して復号化する仕組みを実現したかったため、今回は上記の様なコードにしている。
（上記の様なコードを書くために、cryptographyのパッケージを直に見てみたが、generate_key() のメソッドは下図の様な処理だったので、入れる引数をrandomではなく固定化させてみた）

3. 保護したモデルの使い方イメージ

# 暗号化したモデル情報の txtを読み込む
with open("encrypted_model.txt", "rb") as file:
    encrypted_data = file.read()

# 同じ鍵を作成して復号化
import base64
from cryptography.fernet import Fernet

# 鍵を作成して、base64でエンコードしたものを暗号化して txt形式で出力
passphrase = b'\xc8P\xde@\xf8\x0e\xfe\x81\xf5\xdf\x1b\xafLDe\xdf\xdfT\xd5\xab6\xa1\x9f\xf0\xd49\xb8\xce\x92A\x1c\x85'
key = base64.urlsafe_b64encode(passphrase)

f = Fernet(key)
encoded_data = f.decrypt(encrypted_data)

# PickleデータをBase64でエンコードする
model_info = base64.b64decode(encoded_data)
model_info

# モデルの読み込み
import pickle
model = pickle.loads(model_info)

モデル利用のツールやアプリについて、上記のコードを.pyc形式のコンパイルされた形で提供すれば、最低限の資産流用の防止にはなるのではないか。。。

Discussion