import os
import shutil
import sys
import subprocess
import re
import time
RED = '\033[91m'
GREEN = '\033[92m'
YELLOW = '\033[93m'
BLUE = '\033[94m' 
RESET = '\033[0m'  


INJECT_SCRIPT_NAME="inject_tmp.py"
WIN_FILE_NAME="win"
VERIFY_TIMEOUT=0.5

def init_exp_all(src_path,dest_path):
    print(f"Try to copy exp  ({src_path} -> {dest_path})")
    try:
        for f in os.listdir(src_path):
            if(f[-3:]==".py"):
                src_file_path=os.path.join(src_path,f)
                dest_file_path=os.path.join(dest_path,f)
                shutil.copy(src_file_path,dest_file_path)
                print(f"{GREEN}Copy successful ({src_file_path} -> {dest_file_path}){RESET}")
    except Exception as msg:
        print(f"{RED}{msg}{RESET}")
        print(f"{RED}Exp copy failed{RESET}")
        return -1

def init_exp(elf_path,dest_path):
    elf_path+=".py"
    dest_path=os.path.join(dest_path,'exp.py')
    print(f"Try to copy exp  ({elf_path} -> {dest_path})")
    try:
        shutil.copy(elf_path,dest_path)
        print(f"{GREEN}Copy successful ({elf_path} -> {dest_path}){RESET}")
    except Exception as msg:
        print(f"{RED}{msg}{RESET}")
        print(f"{RED}Exp copy failed{RESET}")
        return -1   
    
def verify_exp(elf_path,exp_name=""):
    tmp_py="" 
    try:
        dir_path=os.path.dirname(elf_path)

        #确定exp的脚本
        if(exp_name == ""):
            script_exp_name=[]
            for f in os.listdir(dir_path):
                if(f[-3:]==".py"):script_exp_name.append(f)
            script_counts=len(script_exp_name)
            assert(script_counts),"There is no python script under the directory"
            if(script_counts==1):
                exp_name=script_exp_name[0]
            else:
                print(f"{YELLOW}There are multiple scripts, please select one{RESET}")
                for i in range(script_counts):
                    print(f"{i} -> {script_exp_name[i]}")
                idx=int(input("idx:"))
                assert(0<=idx<script_counts),"Index wrong"
                exp_name=script_exp_name[idx]
                
        print(f"{BLUE}[*]Found {exp_name} -> Attack ...{RESET}")
    
        exp_path=os.path.join(dir_path,exp_name)
    #注入判断语句
        with open(exp_path,"r") as f:
            script_content=f.read()
        #匹配连接的变量名称
        pattern = r"^(?!#)\s*(\w+).interactive"
        match=re.search(pattern ,script_content,re.MULTILINE)
        assert(match),f"Failed to match the script RE" 
        
        # print(script_content)
        
        PID_virtualname=match.group(1)
        #注入利用win文件判断，先清除win文件
        if( os.path.exists(WIN_FILE_NAME) and os.path.isfile(WIN_FILE_NAME) ):os.remove(WIN_FILE_NAME)

        script_split=script_content.split("\n")
        inject_payload=f"""\n
\t{PID_virtualname}.sendline(b"clear;echo 'Successful Attack {elf_path}' >> {WIN_FILE_NAME};")
\tstrs={PID_virtualname}.recvuntil(b'mowen',timeout={VERIFY_TIMEOUT})
\tif(strs==b''):{PID_virtualname}.close()
except:
\tpass
finally:
\t{PID_virtualname}.close()
        """
        tmp_py=os.path.join(dir_path,INJECT_SCRIPT_NAME)
        with open(tmp_py,"w+") as f:
            f.write("try: \n")
            for s in script_split:
                if("interactive" in s):
                    f.write(inject_payload+"\n")
                    continue
                f.write("\t"+s+"\n")
                
        print(f"Inject payload successful! Start program...")
        #启动elf
        cmd=[elf_path]
        elf_process=subprocess.Popen(cmd,stdout=subprocess.PIPE,stderr=subprocess.PIPE)

        #开始执行exp
        print(f"Run payload...")

        cmd=f"cd {dir_path} ;python3 {INJECT_SCRIPT_NAME};"
        old_time=time.time()
        subprocess.run(cmd,check=True,shell=True,capture_output=True)
        run_time=time.time()-old_time
        print(f"run end process({run_time:.3f})")

        assert(os.path.exists(WIN_FILE_NAME)),f"Failed to attack" 

        with open(WIN_FILE_NAME,"r") as f:
            print(f"{GREEN}{f.read()}{RESET}")
            
        #关闭进程
        elf_process.kill()
        elf_process.wait()
    except AssertionError as msg:
        print(f"{RED}{msg}{RESET}")
        return -1
    except Exception as msg:
        print(f"{RED}{msg}{RESET}")
        return -1
    finally:
        # with open(tmp_py,"r") as f:
        #     print(f.read())
        #删除文件
        if(os.path.exists(WIN_FILE_NAME)):os.remove(WIN_FILE_NAME) 
        if(os.path.exists(tmp_py)):os.remove(tmp_py)

            




if __name__ == "__main__":
    elf_path=sys.argv[1]
    init_exp(os.path.dirname(elf_path),"./")
    verify_exp(elf_path)