import os
import shlex
import re
from collections import OrderedDict
from tools.ToolBase import ToolBase

class HydraTool(ToolBase):
    def validate_instruction(self, instruction):
        timeout = 0
        current_path = os.path.dirname(os.path.realpath(__file__))
        #hydra过滤  需要判断指令中添加字典文件存不存在
        match_p = re.search(r'-P\s+([^\s]+)', instruction)
        match_l = re.search(r'-L\s+([^\s]+)', instruction)
        if match_p:
            str_p = match_p.group(1)
            #判断文件是否存在
            #if not os.path.exists(str_p):   #文件不存在要替换
            new_pass_path = os.path.join(current_path, "../payload", "passwords")
            instruction = instruction.replace(str_p,new_pass_path)
        if match_l:
            str_l = match_l.group(1)
            #判断文件是否存在
            #if not os.path.exists(str_l):
            new_user_path = os.path.join(current_path, "../payload", "users")
            instruction = instruction.replace(str_l, new_user_path)

        #不是双字典的情况加-f
        if "-l" in instruction or "-p" in instruction:
            if "-f" not in instruction:
                instruction = instruction.strip() + " -f"   #当是单密码,或单用户名时,使用成功即停止模式

        #取消-v -V
        instruction = instruction.replace(" -V "," ")
        instruction = instruction.replace(" -v "," ")
        instruction = instruction.replace(" -vV","")
        #加-o   存在个不确定项是:若没有匹配到,输出文件里面是只有一行执行的命令,没有结果描述
        if " -o" not in instruction:
            instruction  = instruction + " -o hydra_result.txt"
        # # 加 -q
        # if " -q" not in instruction:
        #     instruction = instruction + " -q"

        return instruction,timeout

    def merge_info(self,result):
        try:
            # 按行分割输出,保留非空行
            lines = [line.strip() for line in result.splitlines() if line.strip() != ""]
            # 使用有序字典统计相同行的出现次数,保持原始顺序
            counts = OrderedDict()
            for line in lines:
                if line in counts:
                    counts[line] += 1
                else:
                    counts[line] = 1
            # 生成整合后的输出,重复的行后面跟上*次数标记
            output_lines = []
            for line, count in counts.items():
                if count > 1:
                    output_lines.append(f"{line} *{count}")
                else:
                    output_lines.append(line)
            consolidated = "\n".join(output_lines)
            return consolidated
        except Exception as e:
            return result

    def analyze_result(self, result,instruction,stderr,stdout):
        #返回结果
        # result = self.merge_info(result)
        # print(result)
        #加文件后缀了
        lines = result.splitlines()
        if len(lines) == 1:
            result = "没有匹配到成功的结果"
        return result