quickhits.py

#!/usr/bin/python3

import os
import sys
import re
import json
import time
import random
import argparse
import requests
from datetime import datetime
from threading import Thread
from queue import Queue
from functools import partial
from urllib.parse import urlparse
from colored import fg, bg, attr
from multiprocessing.dummy import Pool
from lockfile import LockFile



def banner():
	print("""
                       _      _    _     _ _
            __ _ _   _(_) ___| | _| |__ (_) |_ ___       _ __  _   _
           / _` | | | | |/ __| |/ / '_ \| | __/ __|     | '_ \| | | |
          | (_| | |_| | | (__|   <| | | | | |_\__ \  _  | |_) | |_| |
           \__, |\__,_|_|\___|_|\_\_| |_|_|\__|___/ (_) | .__/ \__, |
              |_|                                       |_|    |___/

                                by @gwendallecoguic

""")
	pass

banner()



startTime = datetime.now()
title_regexp = re.compile('<title>(.*?)</title>', re.IGNORECASE|re.DOTALL)

def testURL( url ):
    time.sleep( 0.3 )
    sys.stdout.write( 'progress: %d/%d\r' %  (t_multiproc['n_current'],t_multiproc['n_total']) )
    t_multiproc['n_current'] = t_multiproc['n_current'] + 1

    if url in t_history and t_history[ url ] == 1:
        return

    if (t_multiproc['n_current'] % 1000) == 0:
        if not lock.is_locked():
            try:
                lock.acquire()
                with open(t_multiproc['f_session'], 'w') as json_file:
                    json.dump(t_history, json_file)
                lock.release()
            except Exception as e:
                a = 1

    t_urlparse = urlparse(url)
    u = t_urlparse.scheme + '_' + t_urlparse.netloc

    if not u in t_exceptions:
        t_exceptions[u] = 0
    if t_exceptions[u] >= 3:
        t_history[ url ] = 1
        # print("skip %s" % t_urlparse.netloc)
        return

    try:
        r = requests.get( url, timeout=5, verify=False, stream=True )
        # if t_multiproc['_grabfiles']:
        #     r = requests.get( url, timeout=2, verify=False, stream=True )
        # else:
        #     r = requests.head( url, timeout=2, verify=False )
    except Exception as e:
        t_exceptions[u] = t_exceptions[u] + 1
        # sys.stdout.write( "%s[-] error occurred: %s%s\n" % (fg('red'),e,attr(0)) )
        return

    t_history[ url ] = 1

    if 'Content-Type' in r.headers:
        content_type = r.headers['Content-Type']
    else:
        content_type = '-'

    # if len(url) > 1:
    #     url = url.strip('_')

    match = title_regexp.search( r.text )
    title = match.group(1).strip() if match else '-'

    ljust = 100
    while ljust < len(url):
        ljust = ljust + 50
    ljust = ljust + 2

    output = '%sS=%d\t\tL=%d\t\tC=%s\t\tT=%s\n' %  (url.ljust(ljust),r.status_code,len(r.text),content_type,title)
    # output = '%sS=%d\t\tL=%d\t\tC=%s\t\tT=%s\n' %  (url.ljust(t_multiproc['u_max_length']),r.status_code,len(r.text),content_type,title)
    # sys.stdout.write( '%s' % output )

    fp = open( t_multiproc['f_output'], 'a+' )
    fp.write( output )
    fp.close()

    if str(r.status_code) in t_codes:
        sys.stdout.write( '%s' % output )

    if t_multiproc['_grabfiles']:
        saveFile( t_multiproc['d_output'], t_urlparse, r )


def saveFile( d_output, t_urlparse, r ):
    # filename = t_urlparse.path.strip('/')
    filename = t_urlparse.path.strip('/') + '_' + t_urlparse.query
    filename = re.sub( '[^0-9a-zA-Z_\-\.]', '_', filename )
    # if len(filename) > 1:
    #     filename = filename.strip('_')

    d_output = d_output +  '/' + t_urlparse.netloc
    f_output = d_output + '/' + filename
    # print(f_output)

    if not os.path.isdir(d_output):
        try:
            os.makedirs( d_output )
        except Exception as e:
            sys.stdout.write( "%s[-] error occurred: %s%s\n" % (fg('red'),e,attr(0)) )
            return

    s_headers = 'HTTP/1.1 ' + str(r.status_code) + ' ' + r.reason + "\n"
    for k,v in r.headers.items():
        s_headers = s_headers + k + ': ' + v + "\n"

    # print(s_headers)
    content = s_headers + "\n" + r.text

    try:
        fp = open( f_output, 'w' )
        fp.write( content )
        fp.close()
    except Exception as e:
        sys.stdout.write( "%s[-] error occurred: %s%s\n" % (fg('red'),e,attr(0)) )
        return


# disable "InsecureRequestWarning: Unverified HTTPS request is being made."
from requests.packages.urllib3.exceptions import InsecureRequestWarning
requests.packages.urllib3.disable_warnings(InsecureRequestWarning)


parser = argparse.ArgumentParser()
parser.add_argument( "-f","--files",help="set file list (required)" )
parser.add_argument( "-g","--no-grab",help="disable file download", action="store_true" )
parser.add_argument( "-o","--hosts",help="set host list (required or -u)" )
# parser.add_argument( "-k","--skip-resolve",help="skip host testing", action="store_true" )
# parser.add_argument( "-r","--redirect",help="follow redirection" )
parser.add_argument( "-s","--no-https",help="disable https", action="store_true" )
parser.add_argument( "-e","--code",help="display only status code separated by comma, default=none" )
parser.add_argument( "-r","--resume",help="resume a previous session by providing the session file" )
parser.add_argument( "-u","--urls",help="set url list (required or -o)" )
parser.add_argument( "-t","--threads",help="threads, default 10" )
parser.parse_args()
args = parser.parse_args()

if args.no_https:
    _https = False
else:
    _https = True

if args.resume:
    resume_file = args.resume
else:
    resume_file = False

if not resume_file:
    t_hosts = []
    if args.hosts:
        if os.path.isfile(args.hosts):
            fp = open( args.hosts, 'r' )
            t_hosts = fp.read().splitlines()
            fp.close()
        else:
            t_hosts.append( args.hosts )
    n_hosts = len(t_hosts)
    sys.stdout.write( '%s[+] %d hosts found: %s%s\n' % (fg('green'),n_hosts,args.urls,attr(0)) )

if not resume_file:
    t_urls = []
    if args.urls:
        if os.path.isfile(args.urls):
            fp = open( args.urls, 'r' )
            t_urls = fp.read().splitlines()
            fp.close()
        else:
            t_urls.append( args.urls )
    n_urls = len(t_urls)
    sys.stdout.write( '%s[+] %d urls found: %s%s\n' % (fg('green'),n_urls,args.urls,attr(0)) )

if not resume_file and n_hosts == 0 and n_urls == 0:
    parser.error( 'hosts/urls list missing' )

if args.files:
    t_files = []
    if os.path.isfile(args.files):
        fp = open( args.files, 'r' )
        t_files = fp.read().splitlines()
        fp.close()
    else:
        t_files.append( args.files )
    n_files = len(t_files)
    sys.stdout.write( '%s[+] %d files found: %s%s\n' % (fg('green'),n_files,args.files,attr(0)) )
elif not resume_file:
    parser.error( 'files list missing' )

if args.no_grab:
    _grabfiles = False
else:
    _grabfiles = True

if args.code:
    t_codes = args.code.split(',')
    t_codes_str = ','.join(t_codes)
else:
    t_codes = []
    t_codes_str = 'none'

if args.threads:
    _threads = int(args.threads)
else:
    _threads = 10


u_max_length = 0
d_output =  os.getcwd()+'/quickhits'
f_output = d_output + '/' + 'output'
if not os.path.isdir(d_output):
    try:
        os.makedirs( d_output )
    except Exception as e:
        sys.stdout.write( "%s[-] error occurred: %s%s\n" % (fg('red'),e,attr(0)) )
        exit()


t_totest = []
t_history = {}

if resume_file:
    if not os.path.isfile(resume_file):
        sys.stdout.write( '%s[-] session file not found: %s%s\n' % (fg('red'),resume_file,attr(0)) )
        exit()
    sys.stdout.write( '%s[+] loading session file: %s%s\n' % (fg('green'),resume_file,attr(0)) )
    f_session = resume_file
else:
    f_session = d_output + '/' + 'sess_' + str(int(time.time()))
    sys.stdout.write( '%s[+] session file: %s%s\n' % (fg('green'),f_session,attr(0)) )

sys.stdout.write( '%s[+] options are -> threads:%d, status_code:%s%s\n' % (fg('green'),_threads,t_codes_str,attr(0)) )

if resume_file:
    sys.stdout.write( '[+] loading datas...\n' )
    with open(resume_file) as json_file:
        t_history = json.load( json_file )
    for url,histo in t_history.items():
        if histo == 0:
            t_totest.append( url )
            l = len(url)
            if l > u_max_length:
                u_max_length = l
    n_totest = len(t_history)
    n_current = len(t_history) - len(t_totest)
else:
    sys.stdout.write( '[+] computing url and host and file list...\n' )

    for url in t_urls:
        for file in t_files:
            u = url.strip().strip('/ ') + '/' + file.strip()
            # u = url.strip().rstrip('/ ') + '/' + file.strip().lstrip('/')
            t_totest.append( u )
            l = len(u)
            if l > u_max_length:
                u_max_length = l

    for host in t_hosts:
        for file in t_files:
            u = 'https' if _https else 'http'
            u = u + '://' + host.strip('/ ') + '/' + file.strip()
            # u = u + '://' + host.strip() + '/' + file.strip().lstrip('/')
            t_totest.append( u )
            l = len(u)
            if l > u_max_length:
                u_max_length = l

    for url in t_totest:
        t_history[ url ] = 0

    n_current = 0
    n_totest = len(t_totest)

sys.stdout.write( '%s[+] %d urls to call%s\n' % (fg('green'),n_totest,attr(0)) )
if n_current:
    sys.stdout.write( '[+] %d already performed\n' % n_current )
sys.stdout.write( '[+] testing...\n' )

# print(n_current)
# print(t_history)
# print(t_totest)
# exit()

random.shuffle(t_totest)
# random.shuffle(t_totest)

t_exceptions = {}
t_multiproc = {
    'n_current': n_current,
    'n_total': n_totest,
    'u_max_length': u_max_length+10,
    'd_output': d_output,
    'f_output': f_output,
    '_grabfiles': _grabfiles,
    'f_session': f_session
}


### OLD METHOD
# pool = Pool( _threads )
# pool.map( testURL, t_totest )
# pool.close()
# pool.join()

# with open(f_session, 'w') as json_file:
#     json.dump(t_history, json_file)

# sys.stdout.write( '%s[+] %d urls called%s\n' % (fg('green'),n_current,attr(0)) )
# print(datetime.now() - startTime)
###


### NEW METHOD, MUCH FASTER
def doWork():
    while True:
        url = q.get()
        testURL( url )
        # doSomethingWithResult(status, url)
        q.task_done()

q = Queue( _threads*2 )
lock = LockFile('./qh')

for i in range(_threads):
    t = Thread( target=doWork )
    t.daemon = True
    t.start()

try:
    for url in t_totest:
        q.put( url )
    q.join()
except KeyboardInterrupt:
    sys.exit(1)

if lock.is_locked():
    lock.release()

with open(f_session, 'w') as json_file:
    json.dump(t_history, json_file)

sys.stdout.write( '%s[+] %d urls called%s\n' % (fg('green'),t_multiproc['n_current'],attr(0)) )
# print(datetime.now() - startTime)