|
5 | 5 |
|
6 | 6 | import shutil |
7 | 7 | import tempfile |
| 8 | +import threading |
| 9 | +import time |
8 | 10 |
|
9 | 11 | WHITELISTED_TAMPERS_PATH = './tampers/whitelisted' |
10 | 12 |
|
@@ -36,23 +38,133 @@ def init_whatwaf(): |
36 | 38 | return tmp_whatwaf_dir |
37 | 39 |
|
38 | 40 |
|
| 41 | +class WhatWafScan(): |
| 42 | + """ threaded class to scan a target """ |
| 43 | + |
| 44 | + def __init__(self, target): |
| 45 | + self.NOT_LAUCHED_STATE = 0 |
| 46 | + self.RUNNING_STATE = 1 |
| 47 | + self.ENDED_STATE = 2 |
| 48 | + |
| 49 | + self.state = self.NOT_LAUCHED_STATE |
| 50 | + self.target = target |
| 51 | + self.waffed_target = None |
| 52 | + self.thread = None |
| 53 | + self.launched = False |
| 54 | + |
| 55 | + def get_waffed_target(self): |
| 56 | + """ return the new target WhatWaf created if the scan ended |
| 57 | + , else return None |
| 58 | + """ |
| 59 | + if self.state is self.ENDED_STATE: |
| 60 | + return self.results |
| 61 | + else: |
| 62 | + return None |
| 63 | + |
| 64 | + def _scan_thread_function(self): |
| 65 | + self.launched = True |
| 66 | + log.debug("Waffing {}".format(self.target.url)) |
| 67 | + self.waffed_target = whatwaf_target(self.target) |
| 68 | + |
| 69 | + def start(self): |
| 70 | + """ scan a target in a new thread and save it in self.waffed_target """ |
| 71 | + self.thread = threading.Thread(target=self._scan_thread_function) |
| 72 | + self.thread.start() |
| 73 | + self.state = self.RUNNING_STATE |
| 74 | + |
| 75 | + def update_status(self): |
| 76 | + """ set self.state to self.ENDED_STATE if self.launched and |
| 77 | + self.thread is dead |
| 78 | + """ |
| 79 | + try: |
| 80 | + if self.launched and not self.thread.is_alive(): |
| 81 | + self.state = self.ENDED_STATE |
| 82 | + except AttributeError: |
| 83 | + pass |
| 84 | + |
| 85 | + def save_target(self): |
| 86 | + """ update the target which resides in the save """ |
| 87 | + save.update_target(self.waffed_target) |
| 88 | + |
| 89 | + def join(self): |
| 90 | + if self.state is self.ENDED_STATE: |
| 91 | + self.thread.join() |
| 92 | + |
| 93 | + def isEnded(self): |
| 94 | + return True if self.state is self.ENDED_STATE else False |
| 95 | + |
| 96 | + def isRunning(self): |
| 97 | + return True if self.state is self.RUNNING_STATE else False |
| 98 | + |
| 99 | + def gotStarted(self): |
| 100 | + return False if self.state is self.NOT_LAUCHED_STATE else True |
| 101 | + |
| 102 | + |
39 | 103 | def wafdetect_stage(args): |
40 | 104 | """init whatwaf with custom tampers and add details to the targets of the |
41 | 105 | save |
42 | 106 | """ |
43 | | - set_whatwaf_path( |
44 | | - init_whatwaf() |
45 | | - ) |
| 107 | + set_whatwaf_path(init_whatwaf()) |
| 108 | + |
| 109 | + targets_queue = [] |
| 110 | + whatwafscan_queue = [] |
| 111 | + |
| 112 | + # add all unwaffed targets to targets_queue |
| 113 | + for target in save.getTargets(): |
| 114 | + if not target.isWaffed(): |
| 115 | + if target is not None: |
| 116 | + targets_queue.append(target) |
| 117 | + |
| 118 | + # create a whatwafscan for every target and add them in whatwafscan_queue |
| 119 | + for target in targets_queue: |
| 120 | + log.debug('Adding {} to the target queue'.format(target.url)) |
| 121 | + whatwafscan_queue.append(WhatWafScan(target)) |
46 | 122 |
|
| 123 | + # constantly check the number of scans launched |
| 124 | + # if it is under 5, launch scan |
| 125 | + # also, check for ended scans. If there are, do scan.save_target(), and |
| 126 | + # remove() it from whatwafscan_queue |
| 127 | + # if there are no scans remaining, break. |
| 128 | + MINIMUM_RUNNING_SCANS = 5 |
47 | 129 | while True: |
48 | | - target = save.getUnwaffedTarget() |
49 | | - if target is not None: |
50 | | - log.debug("Waffing {}".format(target.url)) |
51 | | - target = whatwaf_target(target) |
52 | | - save.updateTarget(target) |
53 | | - else: |
54 | | - log.debug("All targets got waffed !") |
| 130 | + # break if there are no scan remaining |
| 131 | + if len(whatwafscan_queue) == 0: |
55 | 132 | break |
56 | 133 |
|
| 134 | + running_scans = 0 |
| 135 | + for scan in whatwafscan_queue: |
| 136 | + # if the scan is finished, save and remove |
| 137 | + scan.update_status() |
| 138 | + if scan.isEnded(): |
| 139 | + log.debug( |
| 140 | + 'Properly ending scan for {}' |
| 141 | + .format(scan.target.url) |
| 142 | + ) |
| 143 | + |
| 144 | + scan.save_target() |
| 145 | + whatwafscan_queue.remove(scan) |
| 146 | + |
| 147 | + # if the scan is running, increment running_scans |
| 148 | + if scan.isRunning(): |
| 149 | + running_scans += 1 |
| 150 | + |
| 151 | + # if there are not enough running scans, launch some |
| 152 | + to_launch = MINIMUM_RUNNING_SCANS - running_scans |
| 153 | + log.debug('Running scans: {}; Scans to launch: {}'.format( |
| 154 | + running_scans, to_launch |
| 155 | + )) |
| 156 | + |
| 157 | + for scan in whatwafscan_queue: |
| 158 | + if to_launch > 0: |
| 159 | + if not scan.gotStarted(): |
| 160 | + log.debug('Starting scan for {}'.format(scan.target.url)) |
| 161 | + scan.start() |
| 162 | + to_launch -= 1 |
| 163 | + else: |
| 164 | + break |
| 165 | + |
| 166 | + # check every 5 seconds |
| 167 | + time.sleep(5) |
57 | 168 |
|
58 | | -print(init_whatwaf()) |
| 169 | + # hello, code reader. If you like reading code, tell me in the issue |
| 170 | + # tracker ! |
0 commit comments