Simple dirbuster with threading

Question

I've made a simple dirbuster. (I wanted to play with threading.)

The program makes requests to fuzz for files and/or directories on a site.

Featuring

• fuzzing for response codes
• fuzzing for files
• threading

You can critique any and all.

Troubling points

• The way it handles threading
• I think I'm over engineering the arguments handler

Code

import threading
from queue import Queue
from textwrap import dedent
from urllib.parse import urljoin
import os
import sys
import re

import argparse
import requests

NUMBER_OF_THREADS = 5
QUEUE = Queue()


def create_workers(response_codes):
    """creates some threads of workers"""
    for _ in range(NUMBER_OF_THREADS):
        thread = threading.Thread(target=work, args=(response_codes,))
        thread.daemon = True
        thread.start()

def work(response_codes):
    """gets the url from the queue and call print_response_code"""
    while True:
        url = QUEUE.get()
        print_response_code(url, response_codes)
        QUEUE.task_done()

def dir_buster(url, wordlist, file_extensions):
    """puts the urls to fuzz in a queue"""
    for word in read_wordlist(wordlist):
        QUEUE.put(urljoin(url, word))
        if not file_extensions is None:
            for ext in file_extensions:
                QUEUE.put(urljoin(url, f"{word}.{ext}"))
    QUEUE.join()

def read_wordlist(wordlist):
    """yields a word from a \n delimited wordlist"""
    with open(wordlist) as f:
        for line in f:
            yield line.rstrip()

def print_response_code(url, response_codes):
    """gets response code from url and prints if matches a condition"""
    code = requests.head(url).status_code
    if response_codes is None or code in response_codes:
        print(f"[{code}]\t{url}")

def parse_arguments():
    """arguments parser"""
    parser = argparse.ArgumentParser(usage='%(prog)s [options] <url>',
                                     description='Dirbuster by @Ludisposed',
                                     formatter_class=argparse.RawDescriptionHelpFormatter,
                                     epilog=dedent('''Examples:
                                                      python dirbust.py  -c "200, 301,        401" -w wordlist.txt -t 10 -e "html, php   ,py" https://mylocalweb:1234'''))
    parser.add_argument('-c', '--code', type=str, help='HTTP response codes to filter on')
    parser.add_argument('-e', '--extension', type=str, help='Filename extensions you want to fuzz')
    parser.add_argument('-t', '--threads', type=int, help='Number of threads')
    parser.add_argument('-w', '--wordlist', type=str, help='Wordlist you want to fuzz with')
    parser.add_argument('url', type=str, help='Url you want to fuzz')
    args = parser.parse_args()

    try:
        requests.head(args.url).status_code
    except Exception as e:
        print("[!] Url is not responding")
        sys.exit(1)

    if args.wordlist is None or not os.path.isfile(args.wordlist):
        print("[!] Wordlist is not valid")
        sys.exit(1)

    if not args.code is None and re.match(r"^(\s*\d{1,3}\s*)(,\s*\d{1,3}\s*)*$", args.code) is None:
        print("[!] Response codes are not valid, logging all response codes")
        args.code = None

    if not args.code is None:
        args.code = list(map(int, re.sub(r"\s+", "", args.code).split(",")))

    if not args.extension is None and re.match(r"^(\s*[a-z]+\s*)(,\s*[a-z]+\s*)*$", args.extension.lower()) is None:
        print("[!] Extensions are not valid, only searching for directories")
        args.extension = None

    if not args.extension is None:
        args.extension = re.sub(r"\s+", "", args.extension).split(",")

    if not args.threads is None:
        global NUMBER_OF_THREADS
        NUMBER_OF_THREADS = args.threads

    return args.url, args.code, args.wordlist, args.extension

def main(url, response_codes, wordlist, file_extensions):
    """main method that will create workers and adds url to the Queue"""
    create_workers(response_codes)
    dir_buster(url, wordlist, file_extensions)

if __name__ == '__main__':
    main(*parse_arguments())

Answer 1

I do have some points in regards to scope. Having code like:

if not args.threads is None:
    global NUMBER_OF_THREADS
    NUMBER_OF_THREADS = args.threads

means you'll encounter weird issues when values change and you need to track them down. Obviously I'm talking big picture for your future code. It's better to move this into your entry point, as you don't really need to care about the scope of args.threads after spawning the threads, correct? That will also eliminate the LARGE_WORD_IN_CAPITALS which means you're also doing something wrong (defining a global). Something as simple as (untested example) to fix this.

from multiprocessing import Queue, cpu_count
thread_count = args.threads or cpu_count() # or os.cpu_count()

Threading - I'm guessing you're using that to "speed up" things? You might want to use multiprocessing instead.

And regarding all your if statements in regards to the args is None code, introduce default values.

Finally, if you have to do all those if statements in the args check - group them together.

if a and b=this:
    ...
if a and b=that:
    ...
if a and b=whatever:
    ...

can be reduced to,

if a: # (into an inner block)
    if b=this:
        ...
    if b=that:
        ...
    if b=whatever:
        ...

This is much cleaner.

Good luck!