import re
from tools.ToolBase import ToolBase

class GobusterTool(ToolBase):

    def validate_instruction(self, instruction):
        '''
            指令过滤
            1.线程默认调整为-t 5 (没有找到-t 就添加）  ---暂时取消
            2.*medium.txt  替换为*small.txt        --- 暂时取消
            3.-p    静默输出，只输出有用结果
            :param instruction:
            :return:
        '''
        # 定义要修改的参数的正则表达式模式
        # thread_pattern = r'-t\s*\d+'
        # wordlist_pattern = r'-w\s*(/.*?/.*?-medium\.txt)'
        # # 检查是否有 -t 参数，若没有则添加 -t 5
        # if not re.search(thread_pattern, instruction):
        #     instruction += ' -t 5'
        #
        # # 检查 -w 后面的字典文件，若是 *medium.txt 则换成 *small.txt
        # if re.search(wordlist_pattern, instruction):
        #     instruction = re.sub(wordlist_pattern, lambda m: m.group(0).replace('-medium.txt', '-small.txt'),
        #                          instruction)
        timeout = 0
        if "-q" not in instruction:
            instruction += ' -q'
        return  instruction,timeout

    def analyze_result(self, result,instruction,stderr,stdout):
        #指令结果分析 -q后对结果进行提取
        #重新生成个结果，400-5个，401-5个，200所有，其他还不知道有什么结果所有
        result = ""
        i_400 = 0
        i_401 = 0
        lines = stdout.splitlines()
        for line in lines:
            if line:
                badd = False
                if "200" in line:
                    badd = True
                elif "400" in line:
                    if i_400 < 5:   #400有5个页面就可以了
                        badd = True
                        i_400 += 1
                elif "401" in line:
                    if i_401 < 5:
                        badd = True
                        i_401 += 1
                else:   #未知项不太确定,先保留
                    badd = True
                if badd:
                    result +='\n'
                    result += line
        return result

if __name__ == '__main__':
    sub  = GobusterTool()
    gobuster_command = "gobuster dir -u http://haitutech.cn -w /usr/directory-list-2.3-medium.txt -x php,html,zip,bak"
    print(sub.validate_instruction(gobuster_command))