add unBanned response code feature

puncher-pan · Jun 8, 2015 · 9736ea7 · 9736ea7
1 parent 649d3d6
commit 9736ea7
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/core/task.py b/core/task.py
@@ -11,6 +11,7 @@ def __init__(self, number, target, resource, extension):
         self.payload_length = 0
         self.payload_filename = None
         self.banned_response_codes = []
+        self.unbanned_response_codes = []
 
         self.content = None
         self.location = ""
@@ -29,13 +30,20 @@ def get_payload_length(self):
     def set_banned_response_codes(self, banned_codes):
         self.banned_response_codes = banned_codes
 
+    def set_unbanned_response_codes(self, unbanned_codes):
+        self.unbanned_response_codes = unbanned_codes
+
     def set_payload_filename(self, payload_filename):
         self.payload_filename = payload_filename
 
     def set_response_code(self, code):
         self.response_code = str(code)
         if self.response_code in self.banned_response_codes:
             self.valid = False
+        if self.unbanned_response_codes and len(self.unbanned_response_codes)>0:
+            self.valid = False
+            if self.response_code in self.unbanned_response_codes:
+                self.valid = True
 
     def set_location(self, location):
         self.location = location