# -*- coding: utf-8 -*-
import os
import pandas as pd
 
basedir = '/Users/slay/project/python/datahack/data_hacking/sql_injection/data'
filelist = os.listdir(basedir)
df_list = []
 
# 循环读取 basedir下面的内容，文件名为 'legit'的是合法内容，malicious的是 恶意sql语句
for file in filelist:
    df = pd.read_csv(os.path.join(basedir,file), sep='|||', names=['raw_sql'], header=None)
    df['type'] = 'legit' if file.split('.')[0] == 'legit' else 'malicious'
    df_list.append(df)
 
# 将内容放入 dataframe对象
dataframe = pd.concat(df_list, ignore_index=True)
 
dataframe.dropna(inplace=True)
 
# 统计内容
print dataframe['type'].value_counts()
 
# 查看前五个
dataframe.head()

select user from admin;
select hello from admin;

import sqlparse
import string
 
def parse_sql(raw_sql):
    parsed_sql = []
    sql = sqlparse.parse(unicode(raw_sql,'utf-8'))
    for parse in sql:
        for token in parse.tokens:
            if token._get_repr_name() != 'Whitespace':
                    parsed_sql.append(token._get_repr_name())
    return parsed_sql
 
sql_one = parse_sql("select 2 from admin")
sql_two = parse_sql("INSERT INTO Persons VALUES ('Gates', 'Bill', 'Xuanwumen 10', 'Beijing')")
 
print "sql one :%s"%(sql_one)
print "sql two :%s"%(sql_two)

dataframe['parsed_sql'] = dataframe['raw_sql'].map(lambda x:parse_sql(x))
dataframe.head()

dataframe['len'] = dataframe['parsed_sql'].map(lambda x:len(x))
dataframe.head()

%matplotlib inline
import matplotlib.pyplot as plt
dataframe.boxplot('len','type')
plt.ylabel('SQL Statement Length')

import numpy as np
from sklearn.preprocessing import LabelEncoder
from sklearn.feature_extraction.text import CountVectorizer
import string
 
vectorizer = CountVectorizer(min_df=1)
le = LabelEncoder()
 
X = vectorizer.fit_transform(dataframe['parsed_sql'].map(lambda x:string.join(x,' ')))
 
x_len = dataframe.as_matrix(['len']).reshape(X.shape[0],1)
 
x = X.toarray()
 
y = le.fit_transform(dataframe['type'].tolist())
 
print x[:100]
print y[:100]

clf = sklearn.ensemble.RandomForestClassifier(n_estimators=30)
 
scores = sklearn.cross_validation.cross_val_score(clf, x, y, cv=10, n_jobs=4) 
 
print scores

from sklearn.cross_validation import train_test_split
# 将数据分割为 训练数据 和 测试数据，训练数据用于训练模型，测试数据用于测试分类器性能。
X_train, X_test, y_train, y_test, index_train, index_test = train_test_split(x, y, dataframe.index, test_size=0.2)
# 开始训练
clf.fit(X_train, y_train)
# 预测
X_pred = clf.predict(X_test)

%matplotlib inline
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
 
cm = confusion_matrix(X_pred,y_test)
print cm
 
# Show confusion matrix in a separate window
plt.matshow(cm)
plt.title('Confusion matrix')
plt.colorbar()
plt.ylabel('True label')
plt.xlabel('Predicted label')
plt.show()

loss = np.zeros(2)
y_probs = clf.predict_proba(X_test)[:,1]
thres = 0.7 # 用0.7的几率来分类
y_pro = np.zeros(y_probs.shape)
y_pro[y_probs>thres]=1.
cm = confusion_matrix(y_test, y_pro)
print cm

def f(s_x):
    loss = np.zeros(2)
    y_probs = clf.predict_proba(X_test)[:,1]
    thres = s_x # This can be set to whatever you'd like
    y_pro = np.zeros(y_probs.shape)
    y_pro[y_probs>thres]=1.
    cm = confusion_matrix(y_test, y_pro)
    counts = sum(cm)
    count = sum(counts)
    if counts[0]>0:
        loss[0]=float(cm[0,1])/count
    else:
        loss[0]=0.01
    if counts[1]>0:
        loss[1]=float(cm[1,0])/count
    else:
        loss[1]=0.01
    return loss
 
# 0.1 到 0.9 之前的 100个数值
x = np.linspace(0.1,0.9,100)
# x输入f(x)之后得到的结果
y = np.array([f(i) for i in x])
# 可视化
plt.plot(x,y)
plt.show()