cors_scan.py

#!/usr/bin/env python

import json
import sys
import os
import argparse
import threading

from common.common import *
from common.logger import Log
from common.corscheck import CORSCheck

import gevent
from gevent import monkey
monkey.patch_all()
from gevent.pool import Pool
from gevent.queue import Queue
from colorama import init

# Globals
results = []

def banner():
    print(("""%s
   ____ ___  ____  ____   ____    _    _   _ _   _ _____ ____  
  / ___/ _ \|  _ \/ ___| / ___|  / \  | \ | | \ | | ____|  _ \ 
 | |  | | | | |_) \___ \| |     / _ \ |  \| |  \| |  _| | |_) |
 | |__| |_| |  _ < ___) | |___ / ___ \| |\  | |\  | |___|  _ < 
  \____\___/|_| \_\____/ \____/_/   \_\_| \_|_| \_|_____|_| \_\
                                                               %s%s
        # Coded By Jianjun Chen - whucjj@gmail.com%s
    """ % ('\033[91m', '\033[0m', '\033[93m', '\033[0m')))


def parser_error(errmsg):
    banner()
    print(("Usage: python " + sys.argv[0] + " [Options] use -h for help"))
    print(("Error: " + errmsg))
    sys.exit()


def parse_args():
    # parse the arguments
    parser = argparse.ArgumentParser(
        epilog='\tExample: \r\npython ' + sys.argv[0] + " -u google.com")
    parser.error = parser_error
    parser._optionals.title = "OPTIONS"
    parser.add_argument(
        '-u', '--url', help="URL/domain to check it's CORS policy")
    parser.add_argument(
        '-i',
        '--input',
        help='URL/domain list file to check their CORS policy')
    parser.add_argument(
        '-t',
        '--threads',
        help='Number of threads to use for CORS scan',
        type=int,
        default=50)
    parser.add_argument('-o', '--output', help='Save the results to json file')
    parser.add_argument(
        '-v',
        '--verbose',
        help='Enable Verbosity and display results in realtime',
        action='store_true',
        default=False)
    parser.add_argument('-d', '--headers', help='Add headers to the request.', default=None, nargs='*')
    parser.add_argument(
        '-T',
        '--timeout',
        help='Set requests timeout (default 5 sec)',
        type=int,
        default=10)
    parser.add_argument('-p', '--proxy', help='Enable proxy (http or socks5)')
    args = parser.parse_args()
    if not (args.url or args.input):
        parser.error("No url inputed, please add -u or -i option")
    if args.input and not os.path.isfile(args.input):
        parser.error("Input file " + args.input + " not exist.")
    return args


# Synchronize results
c = threading.Condition()

def scan(cfg):
    log = cfg["logger"]
    global results

    while not cfg["queue"].empty():
        try:
            item = cfg["queue"].get(timeout=1.0)
            cors_check = CORSCheck(item, cfg)
            msg = cors_check.check_one_by_one()

            # Keeping results to be written to file only if needed
            if log.filename and msg:
                c.acquire()
                results.append(msg)
                c.release()
        except Exception as e:
            print(e)
            break

""" 
CORScanner library API interface for other projects to use. This will check 
the CORS policy for a given URL. Example Usage:

>>> from CORScanner.cors_scan import cors_check
>>> ret = cors_check("https://www.instagram.com", None)
>>> ret
{'url': 'https://www.instagram.com', 'type': 'reflect_origin',...}

"""
def cors_check(url, headers=None):
    # 0: 'DEBUG', 1: 'INFO', 2: 'WARNING', 3: 'ALERT', 4: 'disable log'
    log = Log(None, print_level=4)
    cfg = {"logger": log, "headers": headers, "timeout": 5}

    cors_check = CORSCheck(url, cfg)
    #msg = cors_check.check_all_in_parallel()
    msg = cors_check.check_one_by_one()
    return msg

def main():
    init()
    args = parse_args()
    #banner()

    queue = Queue()
    log_level = 1 if args.verbose else 2  # 1: INFO, 2: WARNING

    log = Log(args.output, log_level)
    cfg = {"logger": log, "queue": queue, "headers": parse_headers(args.headers),
        "timeout": args.timeout, "proxy": args.proxy}

    read_urls(args.url, args.input, queue)

    sys.stderr.write("Starting CORS scan...(Tips: this may take a while, add -v option to enable debug info)\n")
    sys.stderr.flush()
    threads = [gevent.spawn(scan, cfg) for i in range(args.threads)]

    try:
        gevent.joinall(threads)
    except KeyboardInterrupt as e:
        pass

    # Writing results file if output file has been set
    if log.filename:
        with open(log.filename, 'w') as output_file:
            output_file.write(json.dumps(results, indent=4))
            output_file.close()
    sys.stderr.write("Finished CORS scanning...\n")
    sys.stderr.flush()


if __name__ == '__main__':
    main()