遇到的问题:
在之前的专栏文章中, 已经通过 TensorFlow + RNN + LSTM 将模型训练好, 并做了持久化操作. 为了将这个模型应用到其他场景, 比如 WAF, 需要传给这个模型待检测的报文载荷, 模型输出最可能的一种或多种攻击类型. 为了追求完美的预测效率, 我尝试编译 C++ 版本的 TensorFlow, 但是一直没有编译成功. 组内讨论了一下, 给出一个方案, 直接用 C++ 调用 Python 文件, 先把架构打通, 再去考虑性能问题.
解决的问题:
TensorFlow 或 TFlearn 训练好深度学习的模型, 然后通过 C++ 批量传入报文载荷, Py 模型计算攻击类型, 传回 C++, 做多分类. 后面, 将载荷分别传入不同的 SVM 支持向量机模型, 做 2 分类. 这里一个很关键的问题, 参数如何传递, 请看下一篇专栏文章.
当然, 这个例子跑的 TensorFlow 模型也不好, 因为我没有给加 Dropout 层, 泛化效果一般般, 还有超参数都是我随便填的, 所以还是要看后面的 TFlearn 高级 API 打造的模型玩法~
运行结果:
编译 C++ 程序: g++ waf.cpp -I /usr/include/python2.7 -lpython2.7
运行:./a.out , 结果如下, 和直接运行 Python 文件效果一样.
- [INFO] Python get module succeed.
- [INFO] Get function (main) succeed.
- [[0.94679815 0.03313421 0.02006768]
- [0.06352095 0.8155419 0.12093709]
- [0.08702958 0.5910286 0.3219419 ]]
核心逻辑:
- Py_Initialize();
- // 获取 Py 文件路径
- // 获取 Py 文件: multi_rnn.py
- // 获取 Py 文件中的函数 main
- // 执行 Py 中的 main 函数
- Py_Finalize();
完整代码:
- /**** waf.cpp ******/
- #include "Python.h"
- #include <iostream>
- using namespace std;
- int main()
- {
- Py_Initialize();
- //Py 文件路径
- string path = "/root/ml/study/1book/test";
- string chdir_cmd = string("sys.path.append(\"") + path + "\")";
- const char* cstr_cmd = chdir_cmd.c_str();
- PyRun_SimpleString("import sys");
- PyRun_SimpleString(cstr_cmd);
- //Py 文件
- PyObject* moduleName = PyString_FromString("multi_rnn");
- PyObject* pModule = PyImport_Import(moduleName);
- if (!pModule)
- {
- cout <<"[ERROR] Python get module failed." << endl;
- return 0;
- }
- cout << "[INFO] Python get module succeed." << endl;
- //Py 文件中的函数 main
- PyObject* pv = PyObject_GetAttrString(pModule, "main");
- if (!pv || !PyCallable_Check(pv))
- {
- cout << "[ERROR] Can't find funftion (main)" << endl;
- return 0;
- }
- cout << "[INFO] Get function (main) succeed." << endl;
- // 执行 Py 中的 main 函数
- PyObject_CallObject(pv,NULL);
- Py_Finalize();
- return 0;
- }
- ##### multi_rnn.py
- import sys
- import urllib
- import numpy as np
- import tensorflow as tf
- from tflearn.data_utils import to_categorical, pad_sequences
- from tensorflow.contrib import rnn
- from sklearn.model_selection import train_test_split
- import time
- def elt(line):
- x = []
- for i, c in enumerate(line):
- c = c.lower()
- x.append(ord(c))
- return x
- def load_file(filename,label,ms=[],ns=[]):
- with open(filename) as f:
- for line in f:
- line = line.strip('\n')
- line = urllib.unquote(line)
- m = elt(line)
- if(label == 2):
- n = 2
- elif(label == 1):
- n = 1
- else:
- n = 0
- ms.append(m)
- ns.append(n)
- print(len(ms))
- def load_files(file1,file2,file3):
- xs = []
- ys = []
- load_file(file1,0,xs,ys)
- load_file(file2,1,xs,ys)
- load_file(file3,2,xs,ys)
- return xs,ys
- def one_hot(x,size):
- v = []
- for x1 in x:
- x2 = [0]*size
- x2[x1] = 1
- v.append(x2)
- return v
- def trans(xs,ys):
- x_train, x_test, y_train, y_test=train_test_split( xs,ys, test_size=0.4,random_state=0)
- x_train = pad_sequences(x_train,maxlen=100,value=0.)
- y_train = one_hot(y_train,size=3)
- x_train = np.array(x_train)
- y_train = np.array(y_train)
- x_test = pad_sequences(x_test,maxlen=100,value=0.)
- y_test = one_hot(y_test,size=3)
- x_test = np.array(x_test)
- y_test = np.array(y_test)
- return x_train,y_train,x_test,y_test
- def test():
- v_normal = "hello/world/waf.php"
- x_xss = "/hello/world/waf.php?key=<script>alert(1)</script>"
- x_sql = "/hello/world/waf.php?key=1' and 1=1#" x_sql ="/hello/world/waf.php?key=-2'union select group_concat(Username),2,3 from Person#"
- v_normal = elt(v_normal)
- x_xss = elt(x_xss)
- x_sql = elt(x_sql)
- v_normal = pad_sequences([v_normal],maxlen=100,value=0.)
- x_xss = pad_sequences([x_xss],maxlen=100,value=0.)
- x_sql = pad_sequences([x_sql],maxlen=100,value=0.)
- v_normal = np.array(v_normal)
- x_xss = np.array(x_xss)
- x_sql = np.array(x_sql)
- return v_normal,x_xss,x_sql
- def train(x_train,y_train,x_test,y_test):
- learning_rate = 0.01
- training_steps = 1000
- batch_size = 3
- display_step = 100
- num_input = 10 # url data input (img shape: 10*10)
- timesteps = 10 # time steps
- num_hidden = 100 # hidden layer num of features
- num_classes = 3 # MNIST total classes (0-9 digits)
- X = tf.placeholder(tf.float32,[None,timesteps,num_input],name="X") # X=(-1,10,10)
- Y = tf.placeholder(tf.float32,[None,num_classes]) # Y=(-1,3)
- weights = {
- 'in': tf.Variable(tf.random_normal([num_input,num_hidden])), #W1=(10,100)
- 'out': tf.Variable(tf.random_normal([num_hidden,num_classes]))#W2=(100,3)
- }
- biases = {
- 'in':tf.Variable(tf.constant(0.1,shape = [num_hidden,])),#b1=(100,)
- 'out':tf.Variable(tf.constant(0.1,shape = [num_classes,]))#b2=(3,)
- }
- X = tf.reshape(X,[-1,num_input]) #X=(-1,10)
- X_in = tf.matmul(X,weights['in']) + biases['in'] #X_in= X*W1+b1
- X_in = tf.reshape(X_in,[-1,timesteps,num_hidden]) #X_in= (-1,10,100)
- cell = rnn.BasicLSTMCell(num_hidden)
- init_state = cell.zero_state(batch_size, dtype=tf.float32)
- outputs, final_state = tf.nn.dynamic_rnn(cell, X_in, initial_state=init_state, time_major=False)
- pred = tf.matmul(final_state[1], weights['out']) + biases['out'] #pred=(-1,3)
- pred_y = tf.nn.softmax(pred)
- loss_op = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=pred, labels=Y))
- optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
- train_op = optimizer.minimize(loss_op)
- correct_pred = tf.equal(tf.argmax(pred, 1), tf.argmax(Y, 1))
- accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
- batch = 3
- saver = tf.train.Saver(max_to_keep=1)
- with tf.Session() as sess:
- sess.run(tf.global_variables_initializer())
- for i in range(len(y_train)/batch):
- batch_x=x_train[(i*batch):((i+1)*batch)]
- batch_y=y_train[(i*batch):((i+1)*batch)]
- batch_x = batch_x.reshape((-1, num_input))
- #print(batch_x.shape,batch_y.shape)
- sess.run(train_op, feed_dict={X: batch_x, Y: batch_y})
- if i % display_step == 0:
- loss, acc = sess.run([loss_op, accuracy], feed_dict={X: batch_x,Y: batch_y})
- #print("Step" + str(i) + ", loss =" + \
- # "{:.4f}".format(loss) + ", acc =" + \
- # "{:.3f}".format(acc))
- tf.add_to_collection('pred_network', pred_y)
- saver.save(sess,'./waf.ckpt')
- print("trained ok...")
- accs = 0
- i = 0
- for i in range(len(y_test)/batch):
- batch_x=x_test[(i*batch):((i+1)*batch)]
- batch_y=y_test[(i*batch):((i+1)*batch)]
- batch_x = batch_x.reshape((-1, num_input))
- acc = sess.run(accuracy,feed_dict={X:batch_x,Y:batch_y})
- accs += acc
- print("evaluated ok! acc:",accs/i)
- v_normal,x_xss,x_sql = test()
- x_pred = np.vstack((v_normal,x_xss,x_sql))
- x_pred = x_pred.reshape((-1, num_input))
- forecast = sess.run(pred_y,feed_dict={X:x_pred})
- print(forecast)
- def waf(v_normal,x_xss,x_sql):
- with tf.Session() as sess:
- saver = tf.train.import_meta_graph("./waf.ckpt.meta")
- saver.restore(sess,tf.train.latest_checkpoint("./"))
- graph = tf.get_default_graph()
- X = graph.get_operation_by_name('X').outputs[0]
- pred_y = tf.get_collection('pred_network')[0]
- #graph = tf.get_default_graph()
- #X = graph.get_operation_by_name('X').outputs[0]
- x_pred = np.vstack((v_normal,x_xss,x_sql))
- x_pred = x_pred.reshape((-1,10,10))
- print(sess.run(pred_y,feed_dict={X:x_pred}))
- def main():
- """
- xs,ys = load_files(sys.argv[1],sys.argv[2],sys.argv[3])
- x_train,y_train,x_test,y_test = trans(xs,ys)
- train(x_train,y_train,x_test,y_test)
- """
- v_normal,x_xss,x_sql = test()
- waf(v_normal,x_xss,x_sql)
- if __name__ == "__main__":
- main()
来源: https://juejin.im/entry/5be29cb0f265da61776b72a7