#Ftp
import ftplib
import re
import os
import ipaddress
import subprocess
import tempfile
from tools.ToolBase import ToolBase

class FtpTool(ToolBase):

    def is_ip_domain(self,str):
        # IP 地址校验（支持 IPv4/IPv6）
        try:
            ipaddress.ip_address(str)
            return True
        except ValueError:
            pass


        # 域名格式校验
        domain_pattern = re.compile(
            r'^(?!(https?://|www\.|ftp://))'  # 排除 URL 协议
            r'([a-zA-Z0-9]([a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)'  # 子域名
            r'+[a-zA-Z]{2,63}$'  # 顶级域名（2-63个字母）
        )

        # 总长度校验（域名最大253字符）
        return bool(domain_pattern.match(str)) and len(str) <= 253

    def test_anonymous_ftp_login(self,host, username='anonymous', password='anonymous@example.com'):
        try:
            # 创建 FTP 客户端实例并连接服务器
            ftp = ftplib.FTP(host)
            # 尝试使用匿名凭据登录
            ftp.login(username, password)
            # 登录成功，打印消息
            res = f"匿名登录成功: {host}"
            # 关闭连接
            ftp.quit()
        except ftplib.all_errors as e:
            # 登录失败，打印错误信息
            res = f"匿名登录失败: {host} - {e}"
        return res

    def validate_instruction(self, instruction):
        timeout = 30
        #modified_code = "ftp匿名登录测试"
        #若有put文件，则替换为payload文件
        new_file = "payload/test.txt"
        new_instr = re.sub(r'(put\s+)\S+', r'\1' + new_file, instruction)
        return new_instr,timeout

    def do_worker_subprocess(self,str_instruction,timeout,ext_params):
        output = ""
        stdout = ""
        stderr = ""
        try:
            if timeout == 0:
                result = subprocess.run(str_instruction, shell=True, capture_output=True, text=True)
            elif timeout > 0:
                result = subprocess.run(str_instruction, shell=True, capture_output=True, text=True, timeout=timeout)
            else:
                print("timeout参数错误")
            stderr = result.stderr
            stdout = result.stdout
        except subprocess.TimeoutExpired as e:
            stdout = e.stdout if e.stdout is not None else ""
            stderr = e.stderr if e.stderr is not None else ""
            ext_params.is_user = True  # 对于超时的也需要人工进行确认，是否是预期的超时
        except Exception as e:
            ext_params.is_user = True
            return False, str_instruction, f"执行失败：{str(e)}", "", ext_params  # 执行失败，提交给人工确认指令的正确性

        output = stdout
        if stderr:
            output += stderr
        if isinstance(output, bytes):  # 若是bytes则转成str
            output = output.decode('utf-8', errors='ignore')
        return output

    def do_worker_script(self,str_instruction,timeout,ext_params):
        # 创建临时文件保存输出
        with tempfile.NamedTemporaryFile(delete=False) as tmpfile:
            output_file = tmpfile.name

        # 构建并执行 script 命令
        script_cmd = f"script -c '{str_instruction}' {output_file}"
        try:
            result = subprocess.run(script_cmd, shell=True, text=True,timeout=timeout)
            # 读取输出文件内容
            with open(output_file, 'r') as f:
                output = f.read()
                lines = output.splitlines()
                # 跳过第一行（Script started）和最后一行（Script done）
                ftp_output = lines[1:-1]
                output = '\n'.join(ftp_output)
        except subprocess.TimeoutExpired:
            output = "命令超时返回"
            try:
                with open(output_file, 'r') as f:
                    partial_output = f.read()
                if partial_output:
                    output += f"\n部分输出:\n{partial_output}"
            except FileNotFoundError:
                pass  # 文件可能未创建
        except subprocess.CalledProcessError as e:
            output = f"错误: {e}"
        finally:
            # 删除临时文件
            try:
                os.remove(output_file)
            except FileNotFoundError:
                pass  # 文件可能未创建
        return output


    #对于非sh命令调用的工具，自己实现命令执行的内容
    def execute_instruction(self, instruction_old):
        ext_params = self.create_extparams()
        # 第一步：验证指令合法性
        instruction,time_out = self.validate_instruction(instruction_old)
        if not instruction:
            return False, instruction_old, "该指令暂不执行！","",ext_params
        # 过滤修改后的指令是否需要判重？同样指令再执行结果一致？待定---#？

        # 第二步：执行指令---需要对ftp指令进行区分判断
        pattern = re.compile(r'ftp\s+-n\s+\S+(\s+\d+)?\s+<<\s*EOF')
        match = pattern.search(instruction)
        if bool(match): #如果是 ftp -n 192.168.204.137 <<EOF 开头
            # output = self.do_worker_subprocess(instruction,time_out,ext_params)
            # if not output:
            output = self.do_worker_script(instruction,time_out,ext_params)

        else: #最后使用ftp匿名登陆验证代码
            target = ""
            for str in instruction_old.split():
                if self.is_ip_domain(str):
                    target = str

            if target:
                output = self.test_anonymous_ftp_login(target)
            else:
                output = f"本地程序暂不支持该指令内容"

        # 第三步：分析执行结果
        analysis = self.analyze_result(output,instruction,"","")

        return True, instruction, analysis,output,ext_params

    def analyze_result(self, result,instruction,stderr,stdout):
        #
        return result