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

def Extract_Functions(idat64_path,file):    #提取伪代码
    cmd=f'''{idat64_path} -A -B -S"extract_c.py" {file}
    '''
    save_dir="result"
    #cmd执行命令
    try:
        old_time=time.time()
        if(not os.path.exists(save_dir)):os.mkdir(save_dir)# 结果保存文件夹
        with open("tmp.txt","w") as f:  #用于传参的下下策
            f.write(file)

        subprocess.run(cmd,shell=True,check=True,text=True,capture_output=True) #bash执行idat64命令

        os.remove("tmp.txt")

        new_time=time.time()
        run_tim=new_time-old_time

        print(f"{GREEN}Extract_Functions  prcoess({run_tim:.4f}s){RESET}")

    except Exception as error:
        print(f"{RED}error-> {error}{RESET}")
        with open("my.log","a+") as f:
            f.write(f"Extract_Functions {file}-> {error}")
        exit(-1)
def help():
    help_msg='''python test.py -e idat64_path file <Extract_Functions>
python test.py -h
    '''
    print(f"{BLUE}{help_msg}{RESET}")
def check_argvs(length):
    if(len(sys.argv)<=length):
        help()
        exit(-1)
def main():
    check_argvs(1)
    flags=sys.argv[1]
    if(flags=="-e"):
        check_argvs(3)    
        idat64_path=sys.argv[2]
        file=sys.argv[3]
        Extract_Functions(idat64_path,file)
        
        
        
        


if __name__ == "__main__":
    main()