'''
# SMTP命令注入测试（手动交互）
telnet haitutech.cn 25
EHLO attacker.com
MAIL FROM: <admin@haitutech.cn>
RCPT TO: <target@example.com> NOTIFY=SUCCESS,FAILURE
'''
import socket

from tools.ToolBase import ToolBase

class TelnetTool(ToolBase):
    def validate_instruction(self, instruction):
        #指令过滤
        timeout = 0
        return instruction,timeout

    def analyze_result(self, result,instruction,stderr,stdout):
        #指令结果分析
        return result

    def smtp_injection_test(self,cmd_str: str):
        """
        测试 SMTP 命令注入漏洞。

        参数:
          cmd_str: 多行字符串，其中第一行为类似 "telnet haitutech.cn 25"
                   后续行为具体的 SMTP 命令，例如：

                   EHLO attacker.com
                   MAIL FROM: <admin@haitutech.cn>
                   RCPT TO: <target@example.com> NOTIFY=SUCCESS,FAILURE

        返回:
          如果 RCPT 命令响应以 250 开头，认为可能存在漏洞（返回 True），
          否则认为安全或不支持该注入（返回 False）。
        """
        lines = cmd_str.strip().splitlines()
        if not lines:
            print("未提供命令")
            return ""

        # 解析第一行，格式应为：telnet host port
        parts = lines[0].strip().split()
        if len(parts) < 3:
            print("第一行格式错误，预期格式: 'telnet host port'")
            return ""

        host = parts[1]
        try:
            port = int(parts[2])
        except ValueError:
            print("端口转换失败")
            return ""

        last_response = ""
        if port == 25:
            try:
                # 建立 TCP 连接
                s = socket.create_connection((host, port), timeout=10)
            except Exception as e:
                #print("连接失败：", e)
                return f"连接失败:{e}"

            # 读取 SMTP Banner
            try:
                banner = s.recv(1024).decode('utf-8', errors='ignore')
                print("Banner:", banner.strip())
            except Exception as e:
                #print("读取 Banner 失败：", e)
                s.close()
                return f"读取 Banner 失败：{e}"

            # 从第二行开始发送 SMTP 命令
            for line in lines[1:]:
                cmd = line.strip()
                if not cmd:
                    continue
                print("发送命令:", cmd)
                try:
                    s.send((cmd + "\r\n").encode())
                except Exception as e:
                    print("发送命令失败：", e)
                    s.close()
                    return ""

                try:
                    response = s.recv(1024).decode('utf-8', errors='ignore')
                    print("收到响应:", response.strip())
                    last_response = response.strip()
                except Exception as e:
                    print("读取响应失败：", e)
                    s.close()
                    return f"读取响应失败：{e}"
            s.close()
        else:
            return ""

        # 以 RCPT 命令的响应作为判断依据
        # 注意：对于支持 DSN 的服务器，250 可能是合法的响应，
        # 但在不支持的情况下，若返回 250 则可能说明命令注入导致了不正常的处理。
        if last_response.startswith("250"):
            return "RCPT 命令返回 250，可能存在 SMTP 命令注入风险。"
        else:
            return "RCPT 命令返回非 250 响应，暂未发现注入风险。"

    def execute_instruction(self, instruction_old):
        '''
        执行指令：验证合法性 -> 执行 -> 分析结果
        :param instruction_old:
        :return:
            bool:true-正常返回给大模型，false-结果不返回给大模型
            str:执行的指令
            str:执行指令的结果
        '''
        ext_params = self.create_extparams()
        # 第一步：验证指令合法性
        instruction,time_out = self.validate_instruction(instruction_old)
        if not instruction:
            return False,instruction_old,"该指令暂不执行！","",ext_params
        #过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？

        # 第二步：执行指令
        output = self.smtp_injection_test(instruction)

        # 第三步：分析执行结果
        analysis = self.analyze_result(output,instruction,"","")
        #指令和结果入数据库
        #？
        if not analysis:    #analysis为“” 不提交LLM
            return False,instruction,analysis,output,ext_params
        return True,instruction, analysis,output,ext_params

if __name__ == "__main__":
    pass