From 287afeca047364389cadea084d6ad763f1bcc8b4 Mon Sep 17 00:00:00 2001
From: cxykevin <cxykevin@yeah.net>
Date: Sat, 10 Feb 2024 16:15:59 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=8C=E6=88=90path=E6=B7=BB=E5=8A=A0?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 package/cli/cli/1.bat |  1 +
 src/action.py         |  8 +++++++-
 src/loader.py         | 22 ++++++++++++++--------
 src/makelnk.py        |  6 ++++--
 src/regwrite.py       | 19 +++++++++++++++++++
 5 files changed, 45 insertions(+), 11 deletions(-)
 create mode 100644 package/cli/cli/1.bat
 create mode 100644 src/regwrite.py

diff --git a/package/cli/cli/1.bat b/package/cli/cli/1.bat
new file mode 100644
index 0000000..e19bb96
--- /dev/null
+++ b/package/cli/cli/1.bat
@@ -0,0 +1 @@
+echo helloworld
\ No newline at end of file
diff --git a/src/action.py b/src/action.py
index 60862e4..3f7d550 100644
--- a/src/action.py
+++ b/src/action.py
@@ -29,10 +29,10 @@ import config
 import utils
 import pickle
 import log
-import uuid
 import shutil
 import subprocess
 import makelnk
+import regwrite
 
 
 def force_copy(pkg_path, pkg_name, data_list, froms, tos):
@@ -161,4 +161,10 @@ def do_action(actions_list):
             except:
                 actionlogs = actionlogs + \
                     f"    make shortcut \"{i[1]}\" failed: {repr(exp)}\r\n"
+        elif i[0] == 7:
+            try:
+                regwrite.add_path(i[1])
+            except:
+                actionlogs = actionlogs + \
+                    f"    add path \"{i[1]}\" failed: {repr(exp)}\r\n"
     return actionlogs
diff --git a/src/loader.py b/src/loader.py
index 59ed718..b99549e 100644
--- a/src/loader.py
+++ b/src/loader.py
@@ -253,14 +253,20 @@ def load_package(pkg_name):
                 if i not in file_json["start"]["icon"]:
                     log.warn(f"load moudle [{pkg_name}] warning: " +
                              f"Load icon syntax error (lost \"{i}\"), igrone")
-                if len(file_json["start"]["icon"]["icon"]) < 2 or file_json["start"]["icon"]["icon"][1] != ':':
-                    file_json["start"]["icon"]["icon"] = pkg_path + \
-                        "/"+file_json["start"]["icon"]["icon"]
-                if len(file_json["start"]["icon"]["command"]) < 2 or file_json["start"]["icon"]["command"][1] != ':':
-                    file_json["start"]["icon"]["command"] = pkg_path + \
-                        "/"+file_json["start"]["icon"]["command"]
-                actions["onload"].append(
-                    (6, file_json["start"]["icon"]["command"], file_json["start"]["icon"]["icon"], file_json["start"]["icon"]["name"]))
+            if len(file_json["start"]["icon"]["icon"]) < 2 or file_json["start"]["icon"]["icon"][1] != ':':
+                file_json["start"]["icon"]["icon"] = pkg_path + \
+                    "/"+file_json["start"]["icon"]["icon"]
+            if len(file_json["start"]["icon"]["command"]) < 2 or file_json["start"]["icon"]["command"][1] != ':':
+                file_json["start"]["icon"]["command"] = pkg_path + \
+                    "/"+file_json["start"]["icon"]["command"]
+            actions["onload"].append(
+                (6, file_json["start"]["icon"]["command"], file_json["start"]["icon"]["icon"], file_json["start"]["icon"]["name"]))
+        if "path" in file_json["start"]:
+            for i in file_json["start"]["path"]:
+                if len(i) < 2 or i[1] != ':':
+                    i = pkg_path + \
+                        "/"+i
+                actions["onload"].append((7, i))
 
 
 def load():
diff --git a/src/makelnk.py b/src/makelnk.py
index 986a89f..bd79b9d 100644
--- a/src/makelnk.py
+++ b/src/makelnk.py
@@ -1,3 +1,7 @@
+#####################################
+##       PEinjector/makelnk        ##
+#####################################
+
 import os
 import log
 import pylnk3
@@ -14,8 +18,6 @@ for i in os.listdir(f"{sysdrive}\\Users"):
 
 
 def makelnk(name, exepath, iconpath):
-    print(f"{sysdrive}\\Users\\{username}" +
-          f"\\Desktop\\{name}.lnk")
     log.info(f"make shortcut \"{name}\"")
     pylnk3.for_file(exepath,
                     lnk_name=f"{sysdrive}\\Users\\{username}" +
diff --git a/src/regwrite.py b/src/regwrite.py
new file mode 100644
index 0000000..a957684
--- /dev/null
+++ b/src/regwrite.py
@@ -0,0 +1,19 @@
+#####################################
+##       PEinjector/regwrite       ##
+#####################################
+
+import winreg
+import log
+
+
+def add_path(path):
+    log.info(f"add path \"{path}\"")
+    with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
+                        r"SYSTEM\CurrentControlSet\Control\Session Manager\Environment", 0, winreg.KEY_READ) as key:
+        old_path = winreg.QueryValueEx(key, "Path")[0]
+    if old_path[-1] != ';':
+        old_path = old_path+";"
+    new_path = old_path+path.replace("/", "\\")
+    with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE,
+                        r"SYSTEM\CurrentControlSet\Control\Session Manager\Environment", 0, winreg.KEY_WRITE) as key:
+        winreg.SetValueEx(key, "Path", 0, winreg.REG_EXPAND_SZ, new_path)