Galaxy Lab

focus on information security

Tensorflow + Android应用安全初探

1 .背景

2017年,谷歌发布了一个专门针对移动设备优化的Tensorflow版本——Tensorflow Lite,并于年末开放了开发者预览版。
当前ML,DL系统通常软件服务商的服务端为用户进行服务,而此次Google希望能将一部分相关处理转移到用户的移动设备上,从而减轻服务端的压力。同时,并且一些需要学习的用户数据涉及用户的敏感信息,并不适合通过网络连接发送给服务商,因此利用Tensorflow Lite可以使得在设备本地进行模型的train与inference。

2. 探究目的

从安全角度来看,相比于传统的服务端集中学习,在移动设备上使用Tensorflow或是Tensorflow Lite会将训练模型等信息集成进APK,然后对外发布。那么所使用的训练模型和参数等信息自然就可能被外部获取。
本文将按照官方的例子对在APP内集成Tensorflow Lite进行分步学习,并且观察最后的样例APP中训练模型相关的信息应该如何获取与还原。

3. 应用开发过程

对于Tensorflow Lite的详细开发入门,可以参考项目官方页面。安装Tensorflow的过程就不多说了,本文将应用开发过程大致分为两个阶段:
1 生成TFLITE文件
2 开发APP

3.1 生成TFLITE文件

TFLITE文件就是Tensorflow Lite保存训练模型与参数的文件,它是由Tensorflow的训练模型文件.pb和参数文件.ckpt合成固定后生成的.pb文件进一步优化而成的,会更小且更适合在移动设备上运行。
大致流程图如下:

那我们就使用项目中的cifar10样例来运行一下…咦?报错了?!这可是官方的例子啊?众所周知,Tensorflow现在超级火,更新也超快,所以样例代码中的API有许多已经被修改。根据报错去一个个搜即可,这里前前后后加起来可能有十几个报错处需要修改。

成功运行后,我们发现会在/tmp/下产生训练相关的文件,使用项目中freeze_graph.py这个工具来对目标训练文件进行图参数的固化从而产生.pb文件。同样由于版本问题,我在运行项目中的这个固化脚本时存在一些报错,本人使用了另一个修改脚本

有了.pb文件后,我们可以使用toco工具将其转化为我们所需要的.tflite文件。该工具全称为"The TensorFlow Lite Optimizing Converter",其可以将.pb文件(包含图与参数信息)优化并转化为.tflite文件,从而使其可以高效地在移动设备上运行。
官方对toco的使用方式如下:

bazel run --config=opt tensorflow/contrib/lite/toco:toco -- \
  --input_file=(pwd)/mobilenet_v1_1.0_224/frozen_graph.pb \
  --input_format=TENSORFLOW_GRAPHDEF  --output_format=TFLITE \
  --output_file=/tmp/mobilenet_v1_1.0_224.lite --inference_type=FLOAT \
  --input_type=FLOAT --input_arrays=input \
  --output_arrays=MobilenetV1/Predictions/Reshape_1 --input_shapes=1,224,224,3

为了方便使用该toco工具,可以直接用bazel先编译成elf文件(64位直接点我),然后方便复制到/usr/bin下使其成为系统工具来使用。使用该工具时需要在众多参数中仔细输入正确的输入输出信息,然后就能成功制作出优化后的.tflite文件了。

3.2 开发APP

此处我们直接使用位于tensorflow/tensorflow/contrib/lite/java/demo中的例子。该样例编译后的APP执行效果为识别摄像头拍摄到的画面中的物品,效果如图:

该项目主要有四个java代码文件:
1.AutoFitTextureView.java
2.CameraActivity.java
3.Camera2BasicFragment.java
4.ImageClassifier.java

前两个很简单,CameraActivity.java是一个activity类,其中包含了一个Carama2BasicFragment。AutoFitTextureView则是提供了一个自动适应屏幕的摄像头预览View。
Camera2BasicFrgment.java则是整个APP的主体,其实现了UI界面,并且使用classifyFrame()方法周期性地进行图像识别。

  private Runnable periodicClassify =
      new Runnable() {
        @Override
        public void run() {
          synchronized (lock) {
            if (runClassifier) {
              classifyFrame();
            }
          }
          backgroundHandler.post(periodicClassify);
        }
      };

classifyFrame()的代码中可以看出它是从AutoFitTextureView中提取图片信息,保存为bitmap,并将其通过classifyFrame(bitmap)方法传递给ImageClassifier进行内容识别,然后返回识别的结果。

  private void classifyFrame() {
    if (classifier == null || getActivity() == null || cameraDevice == null) {
      showToast("Uninitialized Classifier or invalid context.");
      return;
    }
    Bitmap bitmap =
        textureView.getBitmap(ImageClassifier.DIM_IMG_SIZE_X, ImageClassifier.DIM_IMG_SIZE_Y);
    String textToShow = classifier.classifyFrame(bitmap);
    bitmap.recycle();
    showToast(textToShow);
  }

下面就需要看ImageClassifier.java了,它也是该应用的核心部分。首先它在构造函数中通过activity来获取asset目录下拥有图和参数信息的.tflite文件和存有全部识别结果标签的label.txt文件(其中tflite文件的内容需要被提取为一个org.tensorflow.lite.Interpreter类的对象)。并且申请了一块缓存imgData用来保存当前需要识别的图片信息。labelProbArray则是用于保存识别结果。

  ImageClassifier(Activity activity) throws IOException {
    tflite = new Interpreter(loadModelFile(activity));
    labelList = loadLabelList(activity);
    imgData =
        ByteBuffer.allocateDirect(
            DIM_BATCH_SIZE * DIM_IMG_SIZE_X * DIM_IMG_SIZE_Y * DIM_PIXEL_SIZE);
    imgData.order(ByteOrder.nativeOrder());
    labelProbArray = new byte[1][labelList.size()];
    Log.d(TAG, "Created a Tensorflow Lite Image Classifier.");
  }

在被Carama2BasicFragment调用的classifyFrame(bitmap)中,首先将传入的bitmap中包含的图片像素色彩信息转移到imgData缓存中,然后调用org.tensorflow.lite.Interpreter的run()来识别imgData中的内容,并将识别结果保存进labelProbArray中,最后用printTopKLabels()来筛选出可能性最高的几项label作为返回值textToShow返回给Carama2BasicFragment进行UI更新展示。

  String classifyFrame(Bitmap bitmap) {
    if (tflite == null) {
      Log.e(TAG, "Image classifier has not been initialized; Skipped.");
      return "Uninitialized Classifier.";
    }
    convertBitmapToByteBuffer(bitmap); 
    // Here's where the magic happens!!!
    long startTime = SystemClock.uptimeMillis();
    tflite.run(imgData, labelProbArray); 
    long endTime = SystemClock.uptimeMillis();
    Log.d(TAG, "Timecost to run model inference: " + Long.toString(endTime - startTime));
    String textToShow = printTopKLabels();
    textToShow = Long.toString(endTime - startTime) + "ms" + textToShow;
    return textToShow;
  }

4 .模型文件的提取与还原

4.1 Tensorflow Lite APP

我们将这个官方样例APP的APK解包,然后在assets目录下找到了该.tflite文件。

由于toco工具并不是单向转化的,它还可以将.tflite文件转化为.pb文件。因此,我们可以再次利用toco工具获得包含图信息与参数信息的.pb文件。运行如下命令:

./toco --input_file=./target_in_your_app.tflite --output_file=test2.pb --input_format=TFLITE --output_format=TENSORFLOW_GRAPHDEF --input_shape=[1,224,224,3] --input_array=conv1 --output_array=softmax_linear

需要注意的是,这些参数对于APP的逆向分析者来说并不是现成的,需要仔细查看反编译的代码,从中找到这些信息。否则toco运行将失败。

将该APK放入JADX等工具中反编译后,在如下部分找到了所需转换的信息。

那么该如何从.pb文件中获取能被我们轻易理解的训练模型信息呢?使用tensorboard即可做到。首先利用如下脚本可以将.pb中的信息提取并保存到指定目录下:

import tensorflow as tf
from tensorflow.python.platform import gfile

with tf.Session() as sess:  
    with tf.gfile.FastGFile("/path/to/the/file.pb", 'rb') as f:  
        graph_def = tf.GraphDef()  
        graph_def.ParseFromString(f.read())  
        tf.import_graph_def(graph_def, name='')  
    writer = tf.summary.FileWriter("/path/you/want", sess.graph)  
    writer.close()  

然后输入命令:

tensorboard --logdir=指定目录

输出中包括一个url,在浏览器中访问这个地址,便可以在相关项目页面中看到该模型的详细信息了。下图便是tensorboard最后显示的模型。

4.2 Tensorflow APP

在/tensorflow/tensorflow/examples/android目录下还有一个样例APP,它没有使用Tensorflow Lite,而是通常Tensorflow。该APP的功能源码此处不再展开分析,直接来看一下其关键部分的代码。

从AndroidManifest.xml中可以看出,该APP安装后会出现四个图标,各自对应四个不同的Activity:

        <activity android:name="org.tensorflow.demo.ClassifierActivity"
                  android:screenOrientation="portrait"
                  android:label="@string/activity_name_classification">
            <intent-filter>
                <action android:name="android.intent.action.MAIN" />
                <category android:name="android.intent.category.LAUNCHER" />
            </intent-filter>
        </activity>

        <activity android:name="org.tensorflow.demo.DetectorActivity"
                  android:screenOrientation="portrait"
                  android:label="@string/activity_name_detection">
            <intent-filter>
                <action android:name="android.intent.action.MAIN" />
                <category android:name="android.intent.category.LAUNCHER" />
            </intent-filter>
        </activity>

        <activity android:name="org.tensorflow.demo.StylizeActivity"
                  android:screenOrientation="portrait"
                  android:label="@string/activity_name_stylize">
            <intent-filter>
                <action android:name="android.intent.action.MAIN" />
                <category android:name="android.intent.category.LAUNCHER" />
            </intent-filter>
        </activity>

        <activity android:name="org.tensorflow.demo.SpeechActivity"
            android:screenOrientation="portrait"
            android:label="@string/activity_name_speech">
            <intent-filter>
                <action android:name="android.intent.action.MAIN" />
                <category android:name="android.intent.category.LAUNCHER" />
            </intent-filter>
        </activity>

进一步分析源码,发现这个APP中的模型文件直接使用了固化参数后的.pb文件。如Stylize功能中使用TensorFlowInferenceInterface来获取stylize_quantized.pb文件中的模型信息进行图像处理。

    private static final String MODEL_FILE = "file:///android_asset/stylize_quantized.pb";
......
    inferenceInterface = new TensorFlowInferenceInterface(getAssets(), MODEL_FILE);
......
    inferenceInterface.feed(
        INPUT_NODE, floatValues, 1, bitmap.getWidth(), bitmap.getHeight(), 3);
    inferenceInterface.feed(STYLE_NODE, styleVals, NUM_STYLES);

    inferenceInterface.run(new String[] {OUTPUT_NODE}, isDebug());
    inferenceInterface.fetch(OUTPUT_NODE, floatValues);
......

在DetectorActivity中更是有三种物体检测模型让我们选择:

  private static final String MB_MODEL_FILE = "file:///android_asset/multibox_model.pb";
  ......
  private static final String TF_OD_API_MODEL_FILE =
      "file:///android_asset/ssd_mobilenet_v1_android_export.pb";
  ......
  private static final String YOLO_MODEL_FILE = "file:///android_asset/graph-tiny-yolo-voc.pb";
  ......

DetectorActivity的效果如下:

由此可知,要想获取这类APP,需要提取其资源文件中的.pb文件。相比于Tensorflow Lite少了一个使用toco工具将.tflite转化回.pb的步骤。有了.pb文件后,便和上文中使用tensorboard的方式一样,对其内部内容进行可视化分析。

5 .总结

从上文中可以看出Tensorflow Lite会将训练模型的图信息与参数信息几乎完全加入到APK中,我们可以从其APK中提取并理解该人工智能模型的实现细节。因此,对于一些涉及敏感商业机密的人工智能模型,厂商在使用Tensorflow+Anroid的开发方案前需要谨慎衡量一下其所可能泄露的技术与带来的风险。
对于希望保护自身知识产权的组织,也可以从上文看出一些可以用来防御的点。第一是通过一些资源保护的手段来使得分析人员无法获取.tflite或.pb文件。第二是当使用Tensorflow Lite时,可以通过一些代码保护技术来使得他人难以获得正确的toco所需参数。
因此,个人认为可以对图参数文件进行加密,或是选用一些不仅仅是保护dex,而且对资源文件也有较高保护技术的加固产品。