#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2017/7/8 15:39
# @Author  : Conan
# @Function: 批量获取存在struts2漏洞的host和其存在的struts2漏洞编号


import os
import sys
import threading

_CURRENT_PATH = os.path.dirname(__file__)
_IP_FILEPATH = os.path.join(_CURRENT_PATH, 'ip_list')
_EXPLOITED_IP_FILEPATH = os.path.join(_CURRENT_PATH, 'exploited_ip_list')
_EXP_DIR_PATH = os.path.join(_CURRENT_PATH, 'struts2_exps')


def get_exp_mods():
    """
    Func: import exp from struts_exps as mod
    Args: None
    Returns: exp dict {'expname1':exp_object1,'expname2',exp_object2}
    """
    exp_dict = {}
    sys.path.append(_EXP_DIR_PATH)
    file_list = os.listdir(_EXP_DIR_PATH)
    for filename in file_list:
        if os.path.splitext(filename)[1] == '.py' and os.path.splitext(filename)[0] != '__init__':
            prefix = os.path.splitext(filename)[0]
            exp_dict[prefix] = __import__(prefix)
    return exp_dict


def get_ip_set_from_file(ip_file):
    """
    Func: get target from the file of ip
    Args: ip file path
    Returns: ip list
    """
    with open(ip_file,'r') as f:
       ip_list = [line.strip() for line in f.readlines()]
    ip_set = set(filter(lambda ip: ip and not ip.startswith("#"), ip_list))
    return ip_set


result_lock = threading.Lock()

def save_result(url,exp_name):
    """
    Func: save result
    Args: url exp_name
    Returns: None
    """
    if result_lock.acquire(True):
        with open(_EXPLOITED_IP_FILEPATH,'a') as f:
            f.writelines(url+' '+exp_name+'\n')
        result_lock.release()

def single_exploit(target,mod_obj):
    '''
    Func: single exploit
    Args: target mod
    Returns: None
    '''
    print '---->try exploiting '+target+' '+'with '+mod_obj.struts2_exp_name+'\n'
    if mod_obj.exploit(target):
        save_result(target, mod_obj.struts2_exp_name)


def batch_exploit():
    """
    Func: try to exploit and save the exploited ip to exploited_ip_list file (ip + struts_number)
    Args: None
    Returns: None
    """
    targets = get_ip_set_from_file(_IP_FILEPATH)
    mods = get_exp_mods()
    threads = []
    for target in targets:
        for name, obj in mods.items():
            threads.append(threading.Thread(target=single_exploit, args=(target, obj), name='try'+target+':'+name))
    for t in threads:
        t.start()
    for t in threads:
        t.join()

if __name__ == "__main__":
    batch_exploit()


