基于强化学习的智能经渗透算法
郝伟 2020/04/14

简介

本算法为基于Q-Learning的强化学习算法的智能化攻击工具模拟演示。通过本算法，可以利用强化学习，对不同Exp工具的渗透率进行学习，从而优化渗透工具的选择过程。

源代码

# -*- coding: utf-8 -*-
"""
@date:  2020/03/31 14:50:44
@author: 郝伟

History
2020/03/31 第1版
2020/04/01 发现问题如下，分析中
           1. 训练的收敛速度不够快
           2. 利用率和权值在训练少时表现不好
           3. 足够训练后的权值在一些情况下仍然不是最优解
           4. 对动态的渗透率的表现仍然未知
2020/04/03 为测试效果调研，发现 plotly 和 dash 效果可以满足需求
2020/04/06 以上的几个问题基本调整优化解决，加入学习曲率系数，以适应动态变化
2020/04/07 最新的效果比较理解，在稳定性上和响应时间上都可以接受
2020/04/09 调整优化代码，以和研发对接
"""

from random import random
import matplotlib.pyplot as plt


def panetrate(e, target=None):
    '''
    模拟的测试函数，使用工具 e 对目标进行攻击，返回结果为布尔型。
    '''
    e['n'] += 1
    return random() < e['rp']

def train(E, n, p = 0.005):
    '''
    训练函数，用于对模型进行训练。参数列表定义如下：
    E: 工具列表，每个工具的定义；
    n: 训练次数；
    p: 学习率，此值越高对变化越敏感，曲线波动越大。
    '''
    for _ in range(n):
        E.sort(key=lambda x:x['lp'], reverse=True)
        for e in E:
            e['lp'] = (1 - p) * e['lp']
            if panetrate(e):
                e['lp'] += p
                break

def show(E, n=10000, orderid='name'):
    '''
    显示工具权值信息。
    '''
    E.sort(key=lambda x:x[orderid])
    print("---------- 训练 {} 次 -----------".format(n))
    print("工具    实际概率   学习概率  总次数")
    for e in E:
        print(" {0}\t{1:.0%}\t{2:6.2%}   {3}".format(
                e['name'], e['rp'], e['lp'], e['n']))
    print("-" * 35)


# 输入工具信息
#     名称， 渗透率，total, p1
E = []
Ecount = 7
for i in range(Ecount):
    # 工具列表，每个工具用字典表示，属性名，值 和 作用 如下所示：
    ei = {'id': int(i),         # 编号
          'name': "e" + str(i), # 名称，默认为 e1, e2, ...
          'rp': 0.1 * i,        # 实际概率（real percentage）
          'n': 0,               # 使用次数
          'lp': 0,              # 学习概率（learning percentage）
          'ex': None}           # 对工具对象的引用
    E.append(ei)

##########################################
#                 主程序
##########################################

# 10次循环换代，每次循环进行1000次模拟测试
loops = 10
trainTimes = 10000

data1 = []
xs = [int(x * trainTimes) for x in range(loops)]
ys = []
for i in range(len(E)):
    ys.append([0] * loops)

e6 = E[6] # 记录最高的一个
for i in range(loops):
    for e in E:
        ys[e['id']][i] = e['lp']
    train(E, trainTimes)
    if i == 4:
        e6['rp'] = 0

show(E, loops * trainTimes)


# 绘制显示部分，由于matplotlib 默认不支持中文，所以用英文表示。
plt.title("RL Test Result")
plt.xlabel("Training Times")
plt.ylabel("Learning Percentage")
labels = ["{}={:4.0%}".format(e['name'], e['rp']) for e in E]
cnames = ['gold', 'red', 'green', 'blue', 'black', 'gray', 'steelblue']
for i in range(len(E)):
    plt.plot(xs,                # x坐标共 loops 个点（默认10个）
             ys[i],             # y 坐标的值，表示学习概率
             color=cnames[i],   # 每个工具对应一种颜色
             linewidth=0.5,     # 线宽
             linestyle="-",     # 线段样式
             #marker = list('.ov1234sp.'),  # 曲线标记样式
             label=labels[i])   # 标记内容
plt.legend(loc='upper left', shadow=True, title="Legend")
plt.show()
≡