webhunt

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import os

import click

from src import echo
from src.component_manager import ComponentManager
from src.component_sniffer import ComponentSniffer
from src.log import setup_logger
from src.utils import confirm_continue

# register main group
main_cmd_group = click.group('main')(lambda: None)


@main_cmd_group.command("scan")
@click.option("-u", "--url", type=click.STRING, help="Target", required=True)
@click.option("-d", "--directory", default=os.path.join(os.getcwd(), "components"), help="Components directory, default ./components")
# request
@click.option("-a", "--aggression", is_flag=True, default=False, help="Open aggression mode")
@click.option("-U", "--user-agent", type=click.STRING, help="Custom user agent")
@click.option("-H", "--header", multiple=True, help="Pass custom header LINE to serve")
@click.option("--disallow-redirect", is_flag=True, default=False, help="Disallow redirect")
# component
@click.option("-c", "--component", multiple=True, help="Specify component")
# max-threads
@click.option("-t", "--max-threads", type=click.INT, default=8, help="Set the maximum number of threads, default 8")
# proxy
@click.option("--proxy", type=click.STRING, help="Set proxy is like: '[HTTP/SOCKS4/SOCKS5]/[username]@[password]/[addr]:[port]' ")
@click.option("--proxy_rdns", is_flag=True, default=False, help="Proxy uses rdns")
# verbose
@click.option("-v", "--verbose", is_flag=True, default=False, help="Output detailed debugging information")
def component_sniffer(url, directory, aggression, user_agent, header, disallow_redirect, component, max_threads, proxy, proxy_rdns, verbose):
    """Component scanning on the target"""
    setup_logger(verbose)

    sniffer = ComponentSniffer(url, directory)
    sniffer.aggression = aggression
    sniffer.max_threads = max_threads
    if header:
        sniffer.headers = header
    if user_agent:
        sniffer.user_agent = user_agent
    if proxy:
        sniffer.set_proxy(proxy, proxy_rdns)
    if disallow_redirect:
        sniffer.allow_redirect = not disallow_redirect

    if component:
        results = sniffer.test(component)
    else:
        results = sniffer.start()
    echo.succ(json.dumps(results, ensure_ascii=False))


@main_cmd_group.command("manage")
@click.option("-d", "--directory", default=os.path.join(os.getcwd(), "components"), help="Components directory, default ./components")
# list
@click.option("-l", "--lists", is_flag=True, default=False, help="List components")
# pull
@click.option("--pull", is_flag=True, default=False, help="Pull components from remote database")
@click.option("--pull_webanalyzer", is_flag=True, default=False, help="Pull components from 'https://github.com/webanalyzer/rules'")
# sync
@click.option("--sync", is_flag=True, default=False, help="Synchronize to remote database")
@click.option("--sync-updating", is_flag=True, default=False, help="Update existing components when synchronizing to remote database")
# database
@click.option("--host", type=click.STRING, default="localhost", help="MySQL database host")
@click.option("--port", type=click.INT, default=3306, help="MySQL database port")
@click.option("--db", type=click.STRING, help="MySQL database name")
@click.option("--user", type=click.STRING, help="MySQL database user")
@click.option("--passwd", type=click.STRING, help="MySQL database password")
# search
@click.option("--search", multiple=True, help="Search component name")
# verbose
@click.option("-v", "--verbose", is_flag=True, default=False, help="Output detailed debugging information")
def component_manager(directory, lists, pull, pull_webanalyzer, sync, host, port, db, user, passwd, sync_updating, search, verbose):
    """Management components"""
    setup_logger(verbose)

    manager = ComponentManager(directory)
    if lists:
        manager.lists()
        return
    elif pull_webanalyzer:
        manager.pull_from_webanalyzer()
    elif pull:
        if not(db and user and passwd):
            echo.fail("Pull component need 'db','user','passwd'.")
            return
        echo.tips(
            "This operation will overwrite all components in the current directory '%s'!" % manager.directory)
        confirm_continue()
        manager.pull_from_remote_database(db, user, passwd, host, port)
    elif sync:
        if not(db and user and passwd):
            echo.fail("Sync component need 'db','user','passwd'.")
            return
        echo.tips("Start sync component info to database...")
        confirm_continue()
        manager.sync(db, user, passwd, host, port, sync_updating)
    elif search:
        manager.search(search)
    else:
        echo.tips("No Action.")


if __name__ == "__main__":
    main_cmd_group()