#####################################
##        PEinjector/hook          ##
#####################################

# This moudle will wait PE loaded all
#     things.

# load moudles
import sys
import os
import time
import pickle
import traceback

# 导入路径
sys.path.append("X:\\PEinjector\\src")
sys.path.append("X:\\PEinjector\\tool")
try:
    os.chdir("X:\\PEinjector\\src")
except:
    sys.path.append(".")
    sys.path.append("..\\tool")
    sys.path.append("src")
    sys.path.append("tool")

# 导入模块
log = __import__("log")
config = __import__("config")
action = __import__("action")
utils = __import__("utils")
regwrite = __import__("regwrite")

# logo
logo = """\033[1m\033[34m
    ___  _____      __        __          
   / _ \\/ __(_)__  /_/__ ____/ /____  ____
  / ___/ _// / _ \\/ / -_) __/ __/ _ \\/ __/
 /_/  /___/_/_//_/ /\\__/\\__/\\__/\\___/_/   
              |___/                        
 _________________________________________
\033[0m\033[?25l"""
print(logo)

# 初始化日志系统
log.init()
log.info("hook start")
try:
    # 等待动画
    sys.stdout.write("\033[3A \033[37m")
    for j in range(5):
        for i in range(41):
            sys.stdout.write(" "+"_"*(i+1)+"\n\033[1A")
            time.sleep(0.02)
        sys.stdout.write("\n\033[1A \033[34m")
        for i in range(41):
            sys.stdout.write(" "+"_"*(i+1)+"\n\033[1A")
            time.sleep(0.02)
        sys.stdout.write("\n\033[1A \033[37m")
    sys.stdout.write("\033[0m\033[?25h\n\n\n")

    # 搜索磁盘和用户名
    sysdrive = os.getenv("SystemDrive")
    username = "Default"
    for i in os.listdir(f"{sysdrive}\\Users"):
        if i not in ("Default", "Default User", "Public", "All Users") and not os.path.isfile(f"{sysdrive}\\Users\\{i}"):
            username = i

    # 等待
    while not os.path.exists(f"{sysdrive}\\Users\\{username}\\Desktop"):
        time.sleep(1)

    # 加载
    log.info("load action \"onload\"")
    with open(config.TEMP_DIR.format(TEMP="X:\\PEinjector")+"\\"+"PEinjector.tmp", "rb") as file:
        lists = pickle.load(file)  # 读取文件
    alog = action.do_action(lists)
    with open(config.ACTIONLOGPATH.format(DISK=utils.find_disk()), "a") as file:
        file.write(alog[0])  # 追加写入日志
    with open(config.TEMP_DIR.format(TEMP="X:\\PEinjector")+"\\"+"remove.tmp", "wb") as file:
        pickle.dump(alog[1], file)  # 写入卸载文件
    # 刷新path
    regwrite.refresh_path()

    # 加载完成
    log.info("done")
except Exception as exp:  # 未知错误
    log.break_err("Exception \n"+str(traceback.format_exc(exp)))
    raise exp
print("This window will close.")
sys.exit(0)