quickhits.py

#!/usr/bin/python3.5

# I don't believe in license.
# You can do whatever you want with this program.

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


startTime = datetime.now()


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 = '-'
    
    output = '%sC=%d\t\tL=%d\t\tT=%s\n' %  (url.ljust(t_multiproc['u_max_length']),r.status_code,len(r.text),content_type)
    # 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
    filename = re.sub( '[^0-9a-zA-Z_\-\.]', '_', filename )
    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
    fp = open( f_output, 'w' )
    fp.write( content )
    fp.close()


# 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)