通过机器学习，识别sql注入 , 从苦B人生到代码实现

楔子

开始之前，先谈一嘴安全问题。不爱看套话请直接看前言。

老话说，没有边界就没有安全，互联网上的所有安全措施，都几乎建立在网络边界上。

原来我们的服务分布,部署方式还很简单。服务部署在内网，统一通过网关对外提供服务。这时只需要在网关建立起安全措施就能很有效的防止攻击。比如大名鼎鼎的ModSecurity，只要配上owasp提供的标准规则，就能有效的防止常见的扫描和注入。

但现在是云计算的时代，容器化的时代。东西向流量潜藏着很大的危险，只要云服务有一个节点被攻破，或者我们集群中的一个服务被入侵，其他节点都将变得不再安全。所以这两年一直在提零信任，云安全。

前言

注入攻击是最常见的网站攻击手段，常年位居安全问题的榜首。而数据库是最核心的服务，所以本文的主要目标是保护数据库，识别sql注入。

我在老东家的时候，有花过很大精力在网关上处理安全问题（waf是自研的）。在那个时候我曾提出过，需要在数据库代理上加一些安全措施，以防止在网关被绕过时，能够进最大限度的保护数据。

可惜当时人手不足，方案也不是很好(基于规则匹配的方案，很花时间)，leader觉得想法不错，但种种原因还是被搁置了。当时是2019年秋天，我还记得那天黄色的树叶稀稀拉拉的点缀在下班的路上，我第一次萌生了辞职的念头。

命运奇妙，东家都换了，我又一次做infra设施的开发，就在这2021年的又一个秋天，这回，我主导做数据库代理层的开发。那必然要把这些问题都设计进去，先加到设计方案里，这次应该不会被pass，希望过阵子我不会被打脸。

当然，时代变了，直接机器学习自动分类，再不用苦哈哈写规则了。效率大大提升，本文就是简单实现一下。

架构

架构简单明了，就是在数据库前面加一层代理，判断sql是否存在风险。当然代理还是主要做数据隔离，服务发现，监控等事情，不在本文详述。

流程

识别sql注入，就是一个分类问题。流程大致如下：

特征获取，就是搞一堆样本，在github上找一些注入的sql，我这里找了一个注入的benchmark，大概有六百多条，具体sql这里就不透露了。然后用chatgpt 生成一堆正常的sql，这里也懒的贴了，有需要可以直接用chatgpt生成。
样本预处理：将这些sql语法分析，分词，替换token。
文本向量化，将上面处理好后的文本进行特征提取，我这里主要用tdidf word2vec bert这三种方式。最终将得到的特征数据分为训练数据和测试数据。
然后训练分类模型，再用测试数据验证准确率。这里主要采用经典的分类器 KNN ridge回归 Logistic回归 随机森林 LGBM

特征处理

加载数据

首先加载数据并对sql进行语法分析，然后替换掉非关键信息

比如一下正常的sql：SELECT book_title, author FROM books WHERE genre = 'Education' ORDER BY id LIMIT 2 -- 注释;

token处理之后变成SELECT FROM WHERE LIMIT ANNOTATION，这里并没有采用传统sql分词方法进行字段变换，经典的做法是将非关键的字替换为一个规定值。我是直接干掉了这些信息。

需要注意的是必须保留sql关键字符，比如这里-- 替换为 ANNOTATION。因为这些都属于sql的关键特征，需要保留。

大概是这样的一个映射表，不全，不同的数据库有很多特殊符号，比如pg的 @>

python 复制代码

symbol_map = {'*':"ALLFIELDS",'--+':"ANNOTATION",'--':"ANNOTATION",'/*':"ANNOTATION", '1':"ONE",'0':"ZERO",'@':"LOCALVAR", '@@':"GLOADVAR", '#':"LOCALTABLE",'##':"GLOADBAR", '||':"CONNECTION",'::':"CONVERTSYM"}

这部分代码就不提供了，都是python的str操作。

最终把文本凭借按行拼接在一个sql_tokens_dataset.txt文档中，大概这样子：

sql_tokens_dataset.txt 复制代码

...
SELECT COUNT FROM HAVING COUNT
SELECT EXTRACT MONTH FROM COUNT ALLFIELDS FROM EXTRACT MONTH FROM
SELECT FROM WHERE DATEDIFF NOW
SELECT MAX FROM
SELECT FROM WHERE
SELECT SUM FROM SUM DESC LIMIT
SELECT COUNT ALLFIELDS FROM
SELECT ALLFIELDS FROM WHERE BETWEEN AND
SELECT ALLFIELDS FROM WHERE
...

然后通过下面的代码加载，并分为训练集和测试集：

python 复制代码

with open("../data/sql_tokens_dataset.txt","r") as f:
    x=[i for i in f]
    #在work2vec时 需要下面这行加载数据
    #x=[i.split() for i in f]
#前622行是注入sql，标记为1，后面是正常的sql 标记为0 ，一共1500多条数据
y=[1 for i in range(622)]
for i in range(622,len(x)):
    y.append(0)
# 分为训练集 和测试
x_train, x_test, y_train, y_test = train_test_split(x,y,test_size=0.2)

tf-idf

tf-idf常被用在搜索系统里，用来统计文本中的关键特征，一篇文章可以因为单词过多，只取前K个，但sql长度有限，并且每个特征都关键，全取。

python 复制代码

from sklearn.feature_extraction.text import TfidfVectorizer

vectorizer=TfidfVectorizer()
x_train = vectorizer.fit_transform(x_train).toarray()
print("x-train n_features,n_sample ---> ",x_train.shape)

x_test = vectorizer.transform(x_test).toarray()
print("x-test n_features,n_sample ---> ",x_test.shape)

work2vec

word2vec是Google研究团队在2013年提出的一种高效训练词向量的模型。

它的模型参数最终呈现为 n*m 矩阵，n是单词数，m是特征数。这些特征的信息用来表述了单词之间的关系。

这就比tf-idf高级了，tf-idf是无法描述词组之间的关系，而word2vec可以，并且随着特征数的增加，这种信息表达的越丰富。一般用300作为超参数。但我们这里显然用不到，只设置了vector_size = 100个。

而特征的构建有两种方法：

CBOW (continuous bag of words): 连续词袋模型，通过上下文来预测当前值。相当于一句话中扣掉一个词，让你猜这个词是什么。
skip-gram : 跳字模型，用当前词来预测上下文。相当于给你一个词，让你猜前面和后面可能出现什么词。

文章往往很大很多，为了加速训练有两种方法负采样和层次softmax。

代码如下：

python 复制代码

from gensim.models import Word2Vec

vector_size = 100 #特征维度
# 设置窗口为5，用默认的cbow
w2v = Word2Vec(x_train, vector_size=vector_size, window=5, min_count=1, workers=4, sg=1)
# 取到特征后，每个词对应一个向量，那么一句话对应一个矩阵。求每列均值最终得到特征向量。
# 我在一篇文章里有看到 最大值效果超过均值 实际测试差不多，可能是我的样本太少导致的。
def average_vec(text):
    vec = np.zeros(vector_size)
    for word in text:
        try:
            vec += w2v.wv[word]
        except KeyError:
            continue
    return vec

x_train = np.array([average_vec(i) for i in x_train])
x_test = np.array([average_vec(i) for i in x_test])

bert

transformer是一种用于序列到序列学习的神经网络模型，主要用于自然语言处理任务，如机器翻译、文本摘要等。它在2017年由 Google 提出，采用了注意力机制来对输入序列进行编码和解码。也就是 Encoder 和 Decoder

而bert可以理解成 input + 若干Encoder + output，encoder可以理解为一个加强版的word2vec模型。而一堆decoder的集合则是gpt。

bert的训练分为两个过程预训练(Pre-training) + 微调（Fine-tuning）

预训练：会根据一定文本提前训练好特征，相当于通用的特征模型
微调：将预训练得到的参数，再根据我们实际的业务场景进行特殊化训练。

预训练一般有两种方法：

掩码语言模型（Masked Language Model，MLM）MLM任务在输入文本中随机掩盖一些标记（通常是15%左右），并要求模型预测这些标记的正确词汇。做完形填空。
NSP任务通过给定两个句子，让模型判断它们是否是连续的语言序列，以此来学习语言的连贯性

使用上还是很简单

python 复制代码

import torch
import numpy as np
import transformers as ppb

model_class, tokenizer_class, pretrained_weights = (ppb.DistilBertModel, ppb.DistilBertTokenizer, 'distilbert-base-uncased')

# 加载模型 这里用的'distilbert-base-uncased' base bert的简化版，更小 更快
tokenizer = tokenizer_class.from_pretrained(pretrained_weights)
model = model_class.from_pretrained(pretrained_weights)

# 将一组文本转为向量
def text_to_vec(ts):
    # 先用tokenizer 分词 ，并且添加special tokens，下面我们要用cls
    x_text_tokens = np.array([tokenizer.encode(i,add_special_tokens=True) for i in ts],dtype=object)
    # 因为句子有长短，需要将x_text_tokens里面的所有行 变成同样长度 用0补齐
    max_len = 0
    for i in x_text_tokens:
        if len(i) > max_len:
            max_len = len(i)
    sql_tokens = np.array([i + [0]*(max_len-len(i)) for i in x_text_tokens])
    #补0的部分需要在attention_mask中标记
    attention_mask = np.where(sql_tokens != 0, 1, 0)

    input_ids = torch.tensor(sql_tokens)  
    attention_mask = torch.tensor(attention_mask)
    with torch.no_grad():
        #使用模型转换，last_hidden_states是多维度的元组
        last_hidden_states = model(input_ids, attention_mask=attention_mask)
    # last_hidden_states[0]是一个三维矩阵： 文本个数*position*768的矩阵
    # 其中position 是上面我们 tokenizer encode 的最大长度，也就是 max_len
    # 我们取第0列，也就是cls，它代表了句子的特征。
    vec = last_hidden_states[0][:,0,:].numpy()
    return vec

# todo 微调，不微调效果也很好，我这里偷个懒。如果微调的话ALBERT效果会更好，那天测一下。

## 生成训练和测试数据，每一行是768长度的特征向量
x_train = text_to_vec(x_train)
x_test = text_to_vec(x_test)

使用模型训练

用常用的分类器看一下效果

python 复制代码

from sklearn.linear_model import LogisticRegression
from sklearn.linear_model import RidgeClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.ensemble import RandomForestClassifier
import lightgbm as lgb

modlist = {}
modlist["logistic回归"] = LogisticRegression(max_iter= 1000)
modlist["岭回归"] = RidgeClassifier(tol=1e-2, solver="sag")
modlist["KNN"] = KNeighborsClassifier(n_neighbors=10)
modlist["随机森林"] = RandomForestClassifier(max_depth=10)
modlist["LGB"] = lgb.LGBMClassifier(verbosity= -1)

对三个个词向量的特征进行测试，打印一些关键信息，

py 复制代码

def benchmark(modlist):
    for k,v in modlist.items():
        fit_start_time = time()
        fit = v.fit(x_train,y_train)
        fit_user_time = time()-fit_start_time
        test_start_time = time()
        accuracy = fit.score(x_test,y_test)
        test_user_time = time()-test_start_time
        print("model:%s fit_use_time:%.6fs test_use_time:%.6fs accuracy:%f" % (k,fit_user_time,test_user_time,accuracy))
benchmark(modlist)

结果

先看tfidf的结果

分类器	训练用时	测试用时	准确率
logistic回归	0.007993s	0.000450s	0.965035
岭回归	0.010776s	0.000454s	0.947552
KNN	0.000493s	0.034211s	0.965035
随机森林	0.123271s	0.010601s	0.961538
LGB	0.056822s	0.001699s	0.972028

word2vec的分类结果

分类器	训练用时	测试用时	准确率
logistic回归	0.010086s	0.000449s	0.916084
岭回归	0.070147s	0.000636s	0.828671
KNN	0.000565s	0.035772s	0.965035
随机森林	0.190193s	0.009779s	0.962517
LGB	0.090480s	0.001629s	0.976014

bert的分类结果：

分类器	训练用时	测试用时	准确率
logistic回归	0.087554s	0.001467s	0.979021
岭回归	0.162863s	0.000987s	0.986014
KNN	0.002441s	0.023777s	0.933566
随机森林	0.331166s	0.010159s	0.975524
LGB	0.668288s	0.002407s	0.979021

我又跑了几次，结果看起来差不多，就偷个懒不多次求均值了，
大概还是能看出来，bert和word2vec好于tfidf，tfidf结果不稳定，样本还是太少，和tokens处理也有关系。
集成学习模型更稳定，效果也好，像logistic这种单模型就很不稳定，但是快

尾语

本文只是简单示例，实际使用还是有很大不同。

比如，对sql检测会导致sql执行性能问题。实际使用时，检测和预处理语句并行执行，只在参数传入后才会拦截，严格场景下，检测可以是后置运行的，避免了性能问题。

再比如，实际用的是大多是线性模型，检测sql为注入sql的概率，在不同程度做不同响应，比如>0.5发警告，>0.8报警 >0.95直接拦截，当然这些参数都是可设置的。

做了两个月，项目顺利上线了，但是sql安全检测模块只做了alpha版，做了但没用，算打脸吗~~~

那个数据库代理我也会做个开源版的出来。