def train_net(self, lr=0.001):
running_loss = 0.0
criterion = nn.BCELoss()
losses = []
optimizer = optim.Adam(self.parameters(), lr)
for i, (payload, website, target) in enumerate(
util.get_website_attacks()
): #TODO FULL payload (DICT!!), website: (header, text), target:(start idx von vuln., länge)
if payload is None or website is None or target is None:
continue
start = time.time()
payload = str(payload)
website_tensor = util.example_to_tensor((website[0] + website[1]))
payload_tensor = util.example_to_tensor(payload)
target = util.generate_target_vuln(website[0], website[1], target)
if is_cuda:
website_tensor = website_tensor.cuda()
payload_tensor = payload_tensor.cuda()
target = target.cuda()
website_var = Variable(website_tensor)
payload_var = Variable(payload_tensor)
output = self(website_var, payload_var) #Size: (1,length,1)
optimizer.zero_grad()
loss = criterion(output, Variable(target))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.data[0]
print('[%5d] loss: %.3f' % (i + 1, loss.data[0]))
losses.append(loss.data[0])
self.save_model()
#print("Epoch took " + str(time.time() - start) + "s to complete")
return losses
def backprop(self):
if len(self.memory) < batch_size:
return -1
x = self.memory.sample(self.batch_size)#[(s, a, n, r)]
batch = tuple(zip(*x))#((s1,s2,...),(a1,a2,...), (n1,n2,...), (r1,r2,...))
assert len(batch[0]) == batch_size and len(batch[3]) == batch_size and len(batch[2]) == batch_size and len(batch[1]) == batch_size
non_final = ByteTensor(tuple(map(lambda s: s is not None, batch[2])))
non_final_next = Variable(torch.cat([util.example_to_tensor(s).cuda() for s in batch[2]]), volatile=True)
state = Variable(torch.cat([util.example_to_tensor(s).cuda() for s in batch[0]]))
state.requires_grad = True
action = Variable(torch.cat([util.char_to_tensor(s).unsqueeze(0).cuda() for s in batch[1]]).type(LongTensor))
reward = Variable(Tensor([s for s in batch[3]]))
model_result = self.model(state)
action_value = self.get_action_value(action, model_result)
#next_value = Variable(torch.zeros(batch_size).cuda().type(FloatTensor))
#prediction = self.model(non_final_next).max(1)[0]
#next_value[non_final] = prediction.cuda()
#next_value.volatile = False
#target_action_value = (next_value * self.gamma) + reward
target_action_value = reward
loss = self.loss_fn(action_value, target_action_value)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.data[0]
def train_net(self, lr=0.001):
"""IF CALLING FROM ANOTHER FILE: USE THE TRAIN FUNCTION OUTSIDE ANY CLASS
train the network with given learning rate"""
running_loss = 0.0
criterion = nn.MSELoss()
losses = []
optimizer = optim.Adam(self.parameters(), lr)
for i, (payload, target, difference) in enumerate(
util.get_website_attacks_differences()):
if payload is None or difference is None or target is None:
continue
start = time.time()
payload = util.padd_payload(str(payload))
website_tensor = util.example_to_tensor(difference)
payload_tensor = util.example_to_tensor(payload)
target = util.generate_target_vuln_fullsite(difference, target)
if is_cuda:
website_tensor = website_tensor.cuda()
payload_tensor = payload_tensor.cuda()
target = target.cuda()
website_var = Variable(website_tensor)
payload_var = Variable(payload_tensor)
output = self(website_var, payload_var)
optimizer.zero_grad()
loss = criterion(output, Variable(target))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.data[0]
print('[%5d] loss: %.3f' % (i + 1, loss.data[0]))
losses.append(loss.data[0])
self.save_model()
#print("Epoch took " + str(time.time() - start) + "s to complete")
return losses
def predict(self, resp, to_inject):
website = util.prepare_headers(resp.headers) + resp.text
payload = str(to_inject)
website_tensor = util.example_to_tensor(website)
payload_tensor = util.example_to_tensor(payload)
if is_cuda:
website_tensor = website_tensor.cuda()
payload_tensor = payload_tensor.cuda()
website_var = Variable(website_tensor)
payload_var = Variable(payload_tensor)
output = self.evaluator(website_var, payload_var) #Size: (1,length,1)
return output
def get_supervised_action(self):
if self.epsilon == -1:
self.epsilon = random.random()
threshold = 0.5 #eps_end + (eps_start - eps_end) * math.exp(-1. * self.done / eps_steps)
self.done = self.done + 1
if self.epsilon > threshold: #TODO
possibles = [s for s in util.get_all_targets_copied() if s.startswith(self.state[1:])]
pick = random.choice(possibles)
#pick = possibles[0] #TODO besser???
if len(self.state)-1 < len(pick):
char = pick[len(self.state)-1:len(self.state)]
else:
#util.get_all_targets_copied().pop(0)
#if len(util.get_all_targets_copied()) == 0:
#util.lines_copied = util.lines(util.targetstrings)#TODO
char = util.get_EOS_token()
self.epsilon = -1
return char
else:
epsi = random.random()
threshold = eps_end + (eps_start - eps_end) * math.exp(-1. * self.done / eps_steps)
self.done = self.done + 1
if epsi > threshold:
state_tensor = util.example_to_tensor(self.state).cuda()
input = Variable(state_tensor, volatile=True)
data = self.model(input).data
char = util.tensor_to_char(data)
#print("picked " + char + " with expected reward of " + str(data.max()))
return char
else:
return random.choice(util.get_letters())
def train_net(self, lr=0.001):
running_loss = 0.0
criterion = nn.BCEWithLogitsLoss()
losses = []
optimizer = optim.Adam(self.parameters(), lr)
for i, item in enumerate(util.get_html()):
start = time.time()
input_tensor = util.example_to_tensor((item[0] + item[1])[:2])
if is_cuda:
input_tensor = input_tensor.cuda()
input_var = Variable(input_tensor)
output = self(input_var)
optimizer.zero_grad()
#exit()
loss = criterion(output, input_var)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.data[0]
if i % 2000 == 1999: # print every 2000 mini-batches
print('[%5d] loss: %.3f' % (i + 1, running_loss / 2000))
losses.append(running_loss / 2000)
running_loss = 0.0
self.save_model()
print("Epoch took " + str(time.time() - start) + " to complete")
return losses
def predict(self, raw_website, attacked_website, payload):
"""IF CALLING FROM ANOTHER FILE: USE THE PREDICT FUNCTION OUTSIDE ANY CLASS
predict if the attack was successful.
provide the attack as full strings of the unattacked website, the attacked website and the payload
"""
diff = util.get_string_difference(raw_website, attacked_website)
if len(diff) == 0:
diff = [" "]
payload = util.padd_payload(payload)
website_tensor = util.example_to_tensor(diff)
payload_tensor = util.example_to_tensor(payload)
if is_cuda:
website_tensor = website_tensor.cuda()
payload_tensor = payload_tensor.cuda()
website_var = Variable(website_tensor)
payload_var = Variable(payload_tensor)
output = self.evaluator(website_var, payload_var)
return output
def train_net_cell(self, lr=0.1):
running_loss = 0.0
criterion = nn.BCELoss()
losses = []
optimizer = optim.Adam(self.parameters(), lr)
list = util.get_all_targets()
random.shuffle(list)
for i, item in enumerate(list):
start = time.time()
out = ""
for i in range(1, len(item)):
target_tensor = util.example_to_tensor(
(item + util.get_EOS_token())[i + 1])
#input_tensor = torch.zeros(target_tensor.size()[0], target_tensor.size()[1], self.hidden_size)
input_tensor = util.example_to_tensor(
item[0:i]) #[item[0] for _ in range(len(item))]
if is_cuda:
target_tensor = target_tensor.cuda()
input_tensor = input_tensor.cuda()
input_var = Variable(input_tensor)
target_var = Variable(target_tensor)
optimizer.zero_grad()
output = self(input_var, len(item))
loss = criterion(output, target_var)
out += util.tensor_to_string(output.data)
#print("output: " + util.tensor_to_string(output.data))
#print("Target: " + util.tensor_to_string(target_var.data))
#print("input: " + util.tensor_to_string(input_var.data))
loss.backward()
torch.nn.utils.clip_grad_norm(self.parameters(), 1)
optimizer.step()
# print statistics
running_loss += loss.data[0]
#print('[%5d] loss: %.3f' % (i + 1, loss.data[0]))
print("cell complete output: " + out)
#losses.append(running_loss)
#running_loss = 0.0
self.save_model()
#print("Epoch took " + str(time.time() - start) + " to complete")
return running_loss / len(list)
def get_unsupervised_action(self):
epsilon = random.random()
threshold = eps_end + (eps_start - eps_end) * math.exp(-1. * self.done / eps_steps)
self.done = self.done + 1
if epsilon > threshold:
state_tensor = util.example_to_tensor(self.state).cuda()
input = Variable(state_tensor, volatile=True)
data = self.model(input).data
char = util.tensor_to_char(data)
#print("picked " + char + " with expected reward of " + str(data.max()))
return char
else:
return random.choice(util.get_letters())