Skip to content

Commit ae17096

Browse files
C-PC-P
C-P
authored and
C-P
committed
Added smartfetch flag
Smartfetch decreases network footprint and provides overall speed increase under common circumstances.
1 parent e574578 commit ae17096

1 file changed

Lines changed: 106 additions & 102 deletions

File tree

‎httpscreenshot.py‎

Lines changed: 106 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,10 @@
99

1010
from selenium import webdriver
1111
from urlparse import urlparse
12+
from random import shuffle
13+
from PIL import Image
14+
from PIL import ImageDraw
15+
from PIL import ImageFont
1216
import multiprocessing
1317
import Queue
1418
import argparse
@@ -19,16 +23,11 @@
1923
import ssl
2024
import M2Crypto
2125
import re
22-
from random import shuffle
2326
import time
24-
# Lines 25-27 are out of date methods as per (http://pillow.readthedocs.org/en/latest/installation.html)
25-
#import Image
26-
#import ImageDraw
27-
#import ImageFont
28-
from PIL import Image
29-
from PIL import ImageDraw
30-
from PIL import ImageFont
3127
import signal
28+
import shutil
29+
import hashlib
30+
3231

3332
reload(sys)
3433
sys.setdefaultencoding("utf8")
@@ -102,7 +101,7 @@ def parseGnmap(inFile, autodetect):
102101
for item in fields:
103102
#Make sure we have an open port with an http type service on it
104103
if (item.find('http') != -1 or autodetect) and re.findall('\d+/open',item):
105-
port = None
104+
port = None
106105
https = False
107106
'''
108107
nmap has a bunch of ways to list HTTP like services, for example:
@@ -144,7 +143,6 @@ def setupBrowserProfile(headless):
144143
print e
145144
time.sleep(1)
146145
continue
147-
148146
return browser
149147

150148

@@ -156,7 +154,7 @@ def writeImage(text, filename, fontsize=40, width=1024, height=200):
156154
image.save(filename)
157155

158156

159-
def worker(urlQueue,tout,debug,headless,doProfile,vhosts,subs,extraHosts,tryGUIOnFail):
157+
def worker(urlQueue, tout, debug, headless, doProfile, vhosts, subs, extraHosts, tryGUIOnFail, smartFetch):
160158
if(debug):
161159
print '[*] Starting worker'
162160

@@ -173,18 +171,20 @@ def worker(urlQueue,tout,debug,headless,doProfile,vhosts,subs,extraHosts,tryGUIO
173171

174172
while True:
175173
#Try to get a URL from the Queue
176-
try:
177-
curUrl = urlQueue.get(timeout=tout)
174+
if urlQueue.qsize() > 0:
175+
try:
176+
curUrl = urlQueue.get(timeout=tout)
177+
except Queue.Empty:
178+
continue
178179
print '[+] '+str(urlQueue.qsize())+' URLs remaining'
179180
screenshotName = urlparse(curUrl[0]).netloc.replace(":", "-")
180181
if(debug):
181182
print '[+] Got URL: '+curUrl[0]
182183
if(os.path.exists(screenshotName+".png")):
183184
if(debug):
184-
print "[-] Screenshot already exists, skipping"
185+
print "[-] Screenshot already exists, skipping"
185186
continue
186-
187-
except Queue.Empty:
187+
else:
188188
if(debug):
189189
print'[-] URL queue is empty, quitting.'
190190
browser.quit()
@@ -204,59 +204,71 @@ def worker(urlQueue,tout,debug,headless,doProfile,vhosts,subs,extraHosts,tryGUIO
204204
writeImage(resp.headers.get('www-authenticate','NO WWW-AUTHENTICATE HEADER'),screenshotName+".png")
205205
continue
206206
elif(resp is not None):
207-
browser.set_window_size(1024, 768)
208-
browser.set_page_load_timeout((tout))
209-
old_url = browser.current_url
210-
browser.get(curUrl[0].strip())
211-
if(browser.current_url == old_url):
212-
print "[-] Error fetching in browser but successfully fetched with Requests: "+curUrl[0]
213-
if(headless):
214-
if(debug):
215-
print "[+] Trying with sslv3 instead of TLS - known phantomjs bug: "+curUrl[0]
216-
browser2 = webdriver.PhantomJS(service_args=['--ignore-ssl-errors=true'], executable_path="phantomjs")
217-
old_url = browser2.current_url
218-
browser2.get(curUrl[0].strip())
219-
if(browser2.current_url == old_url):
207+
208+
resp_hash = hashlib.md5(resp.text).hexdigest()
209+
210+
if smartFetch and resp_hash in hash_basket:
211+
#We have this exact same page already, copy it instead of grabbing it again
212+
print "[+] Pre-fetch matches previously imaged service, no need to do it again!"
213+
shutil.copy2(hash_basket[resp_hash]+".html",screenshotName+".html")
214+
shutil.copy2(hash_basket[resp_hash]+".png",screenshotName+".png")
215+
else:
216+
if smartFetch:
217+
hash_basket[resp_hash] = screenshotName
218+
219+
browser.set_window_size(1024, 768)
220+
browser.set_page_load_timeout((tout))
221+
old_url = browser.current_url
222+
browser.get(curUrl[0].strip())
223+
if(browser.current_url == old_url):
224+
print "[-] Error fetching in browser but successfully fetched with Requests: "+curUrl[0]
225+
if(headless):
220226
if(debug):
221-
print "[-] Didn't work with SSLv3 either..."+curUrl[0]
222-
browser2.close()
223-
else:
224-
print '[+] Saving: '+screenshotName
225-
html_source = browser2.page_source
226-
f = open(screenshotName+".html",'w')
227-
f.write(html_source)
228-
f.close()
229-
browser2.save_screenshot(screenshotName+".png")
230-
browser2.close()
231-
continue
232-
233-
if(tryGUIOnFail and headless):
234-
print "[+] Attempting to fetch with FireFox: "+curUrl[0]
235-
browser2 = setupBrowserProfile(False)
236-
old_url = browser2.current_url
237-
browser2.get(curUrl[0].strip())
238-
if(browser2.current_url == old_url):
239-
print "[-] Error fetching in GUI browser as well..."+curUrl[0]
240-
browser2.close()
241-
continue
227+
print "[+] Trying with sslv3 instead of TLS - known phantomjs bug: "+curUrl[0]
228+
browser2 = webdriver.PhantomJS(service_args=['--ignore-ssl-errors=true'], executable_path="phantomjs")
229+
old_url = browser2.current_url
230+
browser2.get(curUrl[0].strip())
231+
if(browser2.current_url == old_url):
232+
if(debug):
233+
print "[-] Didn't work with SSLv3 either..."+curUrl[0]
234+
browser2.close()
235+
else:
236+
print '[+] Saving: '+screenshotName
237+
html_source = browser2.page_source
238+
f = open(screenshotName+".html",'w')
239+
f.write(html_source)
240+
f.close()
241+
browser2.save_screenshot(screenshotName+".png")
242+
browser2.close()
243+
continue
244+
245+
if(tryGUIOnFail and headless):
246+
print "[+] Attempting to fetch with FireFox: "+curUrl[0]
247+
browser2 = setupBrowserProfile(False)
248+
old_url = browser2.current_url
249+
browser2.get(curUrl[0].strip())
250+
if(browser2.current_url == old_url):
251+
print "[-] Error fetching in GUI browser as well..."+curUrl[0]
252+
browser2.close()
253+
continue
254+
else:
255+
print '[+] Saving: '+screenshotName
256+
html_source = browser2.page_source
257+
f = open(screenshotName+".html",'w')
258+
f.write(html_source)
259+
f.close()
260+
browser2.save_screenshot(screenshotName+".png")
261+
browser2.close()
262+
continue
242263
else:
243-
print '[+] Saving: '+screenshotName
244-
html_source = browser2.page_source
245-
f = open(screenshotName+".html",'w')
246-
f.write(html_source)
247-
f.close()
248-
browser2.save_screenshot(screenshotName+".png")
249-
browser2.close()
250264
continue
251-
else:
252-
continue
253-
254-
print '[+] Saving: '+screenshotName
255-
html_source = browser.page_source
256-
f = open(screenshotName+".html",'w')
257-
f.write(html_source)
258-
f.close()
259-
browser.save_screenshot(screenshotName+".png")
265+
266+
print '[+] Saving: '+screenshotName
267+
html_source = browser.page_source
268+
f = open(screenshotName+".html",'w')
269+
f.write(html_source)
270+
f.close()
271+
browser.save_screenshot(screenshotName+".png")
260272
except Exception as e:
261273
print e
262274
print '[-] Something bad happened with URL: '+curUrl[0]
@@ -273,17 +285,13 @@ def worker(urlQueue,tout,debug,headless,doProfile,vhosts,subs,extraHosts,tryGUIO
273285

274286

275287
def doGet(*args, **kwargs):
276-
url = args[0]
277-
doVhosts = kwargs['vhosts']
278-
urlQueue = kwargs['urlQueue']
279-
subs = kwargs['subs']
280-
extraHosts = kwargs['extraHosts']
281-
del kwargs['extraHosts']
282-
del kwargs['urlQueue']
283-
del kwargs['vhosts']
284-
del kwargs['subs']
285-
286-
kwargs['allow_redirects']=False
288+
url = args[0]
289+
doVhosts = kwargs.pop('vhosts' ,None)
290+
urlQueue = kwargs.pop('urlQueue' ,None)
291+
subs = kwargs.pop('subs' ,None)
292+
extraHosts = kwargs.pop('extraHosts',None)
293+
294+
kwargs['allow_redirects'] = False
287295

288296
resp = requests.get(url[0],**kwargs)
289297

@@ -295,10 +303,11 @@ def doGet(*args, **kwargs):
295303
port = urlparse(url[0]).port
296304
if(port is None):
297305
port = 443
298-
cert = ssl.get_server_certificate((host,port),ssl_version=ssl.PROTOCOL_SSLv23)
299-
x509 = M2Crypto.X509.load_cert_string(cert)
306+
307+
cert = ssl.get_server_certificate((host,port),ssl_version=ssl.PROTOCOL_SSLv23)
308+
x509 = M2Crypto.X509.load_cert_string(cert)
300309
subjText = x509.get_subject().as_text()
301-
names = re.findall("CN=([^\s]+)",subjText)
310+
names = re.findall("CN=([^\s]+)",subjText)
302311

303312
try:
304313
altNames = x509.get_ext('subjectAltName').get_value()
@@ -330,10 +339,7 @@ def doGet(*args, **kwargs):
330339
extraHosts[name] = 1
331340
urlQueue.put(['https://'+name+':'+str(port),False,url[2]])
332341
print '[+] Added host '+name
333-
334-
335342
return resp
336-
337343
else:
338344
return resp
339345

@@ -384,6 +390,10 @@ def sslError(e):
384390
else:
385391
return False
386392

393+
def signal_handler(signal, frame):
394+
print "[-] Ctrl-C received! Killing Thread(s)..."
395+
sys.exit(0)
396+
signal.signal(signal.SIGINT, signal_handler)
387397

388398
if __name__ == '__main__':
389399
parser = argparse.ArgumentParser()
@@ -399,6 +409,7 @@ def sslError(e):
399409
parser.add_argument("-dB","--dns_brute",help='Specify a DNS subdomain wordlist for bruteforcing on wildcard SSL certs')
400410
parser.add_argument("-r","--retries",type=int,default=0,help='Number of retries if a URL fails or timesout')
401411
parser.add_argument("-tG","--trygui",action='store_true',default=False,help='Try to fetch the page with FireFox when headless fails')
412+
parser.add_argument("-sF","--smartfetch",action='store_true',default=False,help='Enables smart fetching to reduce network traffic, also increases speed if certain conditions are met.')
402413

403414
args = parser.parse_args()
404415

@@ -413,14 +424,12 @@ def sslError(e):
413424
urls = []
414425
for host,ports in hosts.items():
415426
for port in ports:
416-
url=''
427+
url = ''
417428
if port[1] == True:
418429
url = ['https://'+host+':'+port[0],args.vhosts,args.retries]
419430
else:
420431
url = ['http://'+host+':'+port[0],args.vhosts,args.retries]
421432
urls.append(url)
422-
423-
424433
else:
425434
print 'Invalid input file - must be Nmap GNMAP'
426435

@@ -437,32 +446,27 @@ def sslError(e):
437446

438447
#shuffle the url list
439448
shuffle(urls)
449+
440450
#read in the subdomain bruteforce list if specificed
441451
subs = []
442452
if(args.dns_brute != None):
443453
subs = open(args.dns_brute,'r').readlines()
454+
444455
#Fire up the workers
445-
urlQueue = multiprocessing.Queue()
446-
manager = multiprocessing.Manager()
447-
hostsDict = manager.dict()
448-
workers = []
456+
urlQueue = multiprocessing.Queue()
457+
manager = multiprocessing.Manager()
458+
hostsDict = manager.dict()
459+
workers = []
460+
hash_basket = {}
449461

450462
for i in range(args.workers):
451-
p = multiprocessing.Process(target=worker,
452-
args=(urlQueue,args.timeout,args.verbose,args.headless,args.autodetect, args.vhosts,subs,hostsDict,args.trygui))
453-
454-
463+
p = multiprocessing.Process(target=worker, args=(urlQueue, args.timeout, args.verbose, args.headless, args.autodetect, args.vhosts, subs, hostsDict, args.trygui, args.smartfetch))
455464
workers.append(p)
456465
p.start()
457-
466+
458467
for url in urls:
459468
urlQueue.put(url)
460469

461470
for p in workers:
462-
try:
463-
p.join()
464-
except KeyboardInterrupt:
465-
print "[-] Ctrl-C received! Sending kill to threads..."
466-
kill_received = True
467-
for p in workers:
468-
p.terminate()
471+
p.join(.001)
472+

0 commit comments

Comments
 (0)