Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 29f502f

Browse files
stamparmstamparm
committed
some refactoring
1 parent 37d78ff commit 29f502f

8 files changed

Lines changed: 25 additions & 47 deletions

File tree

‎lib/core/common.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1773,15 +1773,15 @@ def readCachedFileContent(filename, mode='rb'):
17731773
"""
17741774

17751775
if filename not in kb.cache.content:
1776-
kb.locks.cacheLock.acquire()
1776+
kb.locks.cache.acquire()
17771777

17781778
if filename not in kb.cache.content:
17791779
checkFile(filename)
17801780
with codecs.open(filename, mode, UNICODE_ENCODING) as f:
17811781
content = f.read()
17821782
kb.cache.content[filename] = content
17831783

1784-
kb.locks.cacheLock.release()
1784+
kb.locks.cache.release()
17851785

17861786
return kb.cache.content[filename]
17871787

@@ -2241,13 +2241,13 @@ def logHTTPTraffic(requestLogMsg, responseLogMsg):
22412241
if not conf.trafficFile:
22422242
return
22432243

2244-
kb.locks.logLock.acquire()
2244+
kb.locks.log.acquire()
22452245

22462246
dataToTrafficFile("%s%s" % (requestLogMsg, os.linesep))
22472247
dataToTrafficFile("%s%s" % (responseLogMsg, os.linesep))
22482248
dataToTrafficFile("%s%s%s%s" % (os.linesep, 76 * '#', os.linesep, os.linesep))
22492249

2250-
kb.locks.logLock.release()
2250+
kb.locks.log.release()
22512251

22522252
def getPageTemplate(payload, place):
22532253
"""

‎lib/core/option.py‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1437,10 +1437,8 @@ def __setKnowledgeBaseAttributes(flushAll=True):
14371437
kb.lastParserStatus = None
14381438

14391439
kb.locks = AttribDict()
1440-
kb.locks.cacheLock = threading.Lock()
1441-
kb.locks.logLock = threading.Lock()
1442-
kb.locks.ioLock = threading.Lock()
1443-
kb.locks.countLock = threading.Lock()
1440+
for _ in ("cache", "count", "index", "io", "limits", "log", "outputs", "value"):
1441+
kb.locks[_] = threading.Lock()
14441442

14451443
kb.matchRatio = None
14461444
kb.multiThreadMode = False

‎lib/core/threads.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,10 @@ def runThreads(numThreads, threadFunction, cleanupFunction=None, forwardExceptio
188188
kb.threadContinue = True
189189
kb.threadException = False
190190

191+
for lock in kb.locks.values():
192+
if lock.locked_lock():
193+
lock.release()
194+
191195
if conf.get("hashDB", None):
192196
conf.hashDB.flush(True)
193197

‎lib/techniques/blind/inference.py‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -320,25 +320,21 @@ def etaProgressUpdate(charTime, index):
320320
threadData.shared.value = [ None ] * length
321321
threadData.shared.index = [ firstChar ] # As list for python nested function scoping
322322

323-
lockNames = ('iolock', 'idxlock', 'valuelock')
324-
for lock in lockNames:
325-
kb.locks[lock] = threading.Lock()
326-
327323
try:
328324
def blindThread():
329325
threadData = getCurrentThreadData()
330326

331327
while kb.threadContinue:
332-
kb.locks.idxlock.acquire()
328+
kb.locks.index.acquire()
333329

334330
if threadData.shared.index[0] >= length:
335-
kb.locks.idxlock.release()
331+
kb.locks.index.release()
336332

337333
return
338334

339335
threadData.shared.index[0] += 1
340336
curidx = threadData.shared.index[0]
341-
kb.locks.idxlock.release()
337+
kb.locks.index.release()
342338

343339
if kb.threadContinue:
344340
charStart = time.time()
@@ -348,10 +344,10 @@ def blindThread():
348344
else:
349345
break
350346

351-
kb.locks.valuelock.acquire()
347+
kb.locks.value.acquire()
352348
threadData.shared.value[curidx-1] = val
353349
currentValue = list(threadData.shared.value)
354-
kb.locks.valuelock.release()
350+
kb.locks.value.release()
355351

356352
if kb.threadContinue:
357353
if showEta:
@@ -388,9 +384,7 @@ def blindThread():
388384
status = ' %d/%d (%d%s)' % (count, length, round(100.0*count/length), '%')
389385
output += status if count != length else " "*len(status)
390386

391-
kb.locks.iolock.acquire()
392387
dataToStdout("\r[%s] [INFO] retrieved: %s" % (time.strftime("%X"), filterControlChars(output)))
393-
kb.locks.iolock.release()
394388

395389
if not kb.threadContinue:
396390
if int(threading.currentThread().getName()) == numThreads - 1:

‎lib/techniques/brute/use.py‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -81,13 +81,13 @@ def tableExistsThread():
8181
threadData = getCurrentThreadData()
8282

8383
while kb.threadContinue:
84-
kb.locks.countLock.acquire()
84+
kb.locks.count.acquire()
8585
if threadData.shared.count < threadData.shared.limit:
8686
table = safeSQLIdentificatorNaming(tables[threadData.shared.count], True)
8787
threadData.shared.count += 1
88-
kb.locks.countLock.release()
88+
kb.locks.count.release()
8989
else:
90-
kb.locks.countLock.release()
90+
kb.locks.count.release()
9191
break
9292

9393
if conf.db and METADB_SUFFIX not in conf.db:
@@ -97,7 +97,7 @@ def tableExistsThread():
9797

9898
result = inject.checkBooleanExpression("%s" % safeStringFormat(BRUTE_TABLE_EXISTS_TEMPLATE, (randomInt(1), fullTableName)))
9999

100-
kb.locks.ioLock.acquire()
100+
kb.locks.io.acquire()
101101

102102
if result and table.lower() not in threadData.shared.unique:
103103
threadData.shared.outputs.append(table)
@@ -112,7 +112,7 @@ def tableExistsThread():
112112
status = '%d/%d items (%d%s)' % (threadData.shared.count, threadData.shared.limit, round(100.0*threadData.shared.count/threadData.shared.limit), '%')
113113
dataToStdout("\r[%s] [INFO] tried %s" % (time.strftime("%X"), status), True)
114114

115-
kb.locks.ioLock.release()
115+
kb.locks.io.release()
116116

117117
try:
118118
runThreads(conf.threads, tableExistsThread, threadChoice=True)
@@ -180,18 +180,18 @@ def columnExistsThread():
180180
threadData = getCurrentThreadData()
181181

182182
while kb.threadContinue:
183-
kb.locks.countLock.acquire()
183+
kb.locks.count.acquire()
184184
if threadData.shared.count < threadData.shared.limit:
185185
column = safeSQLIdentificatorNaming(columns[threadData.shared.count])
186186
threadData.shared.count += 1
187-
kb.locks.countLock.release()
187+
kb.locks.count.release()
188188
else:
189-
kb.locks.countLock.release()
189+
kb.locks.count.release()
190190
break
191191

192192
result = inject.checkBooleanExpression(safeStringFormat(BRUTE_COLUMN_EXISTS_TEMPLATE, (column, table)))
193193

194-
kb.locks.ioLock.acquire()
194+
kb.locks.io.acquire()
195195

196196
if result:
197197
threadData.shared.outputs.append(column)
@@ -205,7 +205,7 @@ def columnExistsThread():
205205
status = '%d/%d items (%d%s)' % (threadData.shared.count, threadData.shared.limit, round(100.0*threadData.shared.count/threadData.shared.limit), '%')
206206
dataToStdout("\r[%s] [INFO] tried %s" % (time.strftime("%X"), status), True)
207207

208-
kb.locks.ioLock.release()
208+
kb.locks.io.release()
209209

210210
try:
211211
runThreads(conf.threads, columnExistsThread, threadChoice=True)

‎lib/techniques/error/use.py‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -173,9 +173,7 @@ def __errorFields(expression, expressionFields, expressionFieldsList, expected=N
173173
return None
174174

175175
if output is not None:
176-
kb.locks.ioLock.acquire()
177176
dataToStdout("[%s] [INFO] %s: %s\r\n" % (time.strftime("%X"), "resumed" if threadData.resumed else "retrieved", safecharencode(output)))
178-
kb.locks.ioLock.release()
179177

180178
if isinstance(num, int):
181179
expression = origExpr
@@ -347,10 +345,6 @@ def errorUse(expression, expected=None, resumeValue=True, dump=False):
347345
debugMsg += "large number of rows. It might take too long"
348346
logger.debug(debugMsg)
349347

350-
lockNames = ('limits', 'outputs')
351-
for lock in lockNames:
352-
kb.locks[lock] = threading.Lock()
353-
354348
try:
355349
def errorThread():
356350
threadData = getCurrentThreadData()

‎lib/techniques/union/use.py‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -278,10 +278,6 @@ def unionUse(expression, unpack=True, dump=False):
278278
debugMsg += "large number of rows. It might take too long"
279279
logger.debug(debugMsg)
280280

281-
lockNames = ('limits', 'value')
282-
for lock in lockNames:
283-
kb.locks[lock] = threading.Lock()
284-
285281
try:
286282
def unionThread():
287283
threadData = getCurrentThreadData()
@@ -326,9 +322,7 @@ def unionThread():
326322
if len(status) > width:
327323
status = "%s..." % status[:width - 3]
328324

329-
kb.locks.ioLock.acquire()
330325
dataToStdout(status, True)
331-
kb.locks.ioLock.release()
332326

333327
runThreads(numThreads, unionThread)
334328

‎lib/utils/crawler.py‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,6 @@ def getTargetUrls(self):
3939
threadData = getCurrentThreadData()
4040
threadData.shared.outputs = oset()
4141

42-
lockNames = ('limits', 'outputs', 'ioLock')
43-
for lock in lockNames:
44-
kb.locks[lock] = threading.Lock()
45-
4642
def crawlThread():
4743
threadData = getCurrentThreadData()
4844

@@ -100,11 +96,9 @@ def crawlThread():
10096
findPageForms(content, current, False, True)
10197

10298
if conf.verbose in (1, 2):
103-
kb.locks.ioLock.acquire()
10499
threadData.shared.count += 1
105100
status = '%d/%d links visited (%d%s)' % (threadData.shared.count, threadData.shared.length, round(100.0*threadData.shared.count/threadData.shared.length), '%')
106101
dataToStdout("\r[%s] [INFO] %s" % (time.strftime("%X"), status), True)
107-
kb.locks.ioLock.release()
108102

109103
threadData.shared.deeper = set()
110104
threadData.shared.unprocessed = set([conf.url])

0 commit comments

Comments
 (0)