From 75af8db0a598f995f1139fbceed6726ba609efa5 Mon Sep 17 00:00:00 2001 From: lifei6671 Date: Sat, 11 Apr 2026 10:38:09 +0800 Subject: [PATCH] Refactor: Improve Edge process management and JSON handling in patch functions --- .codex | 0 main.go | 196 ++++++++++++++++++++++++++++++------------------ main_test.go | 205 +++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 328 insertions(+), 73 deletions(-) create mode 100644 .codex create mode 100644 main_test.go diff --git a/.codex b/.codex new file mode 100644 index 0000000..e69de29 diff --git a/main.go b/main.go index 65f0931..1b42cc4 100644 --- a/main.go +++ b/main.go @@ -17,11 +17,12 @@ Patch Edge Copilot - Go版本 */ package main - import ( "encoding/json" "fmt" + "io" "os" + "os/exec" "path/filepath" "runtime" "strings" @@ -29,18 +30,6 @@ import ( "github.com/shirou/gopsutil/v3/process" ) -type LocalState struct { - VariationsCountry string `json:"variations_country"` -} - -type Preferences struct { - Browser Browser `json:"browser"` -} - -type Browser struct { - ChatIPEligibilityStatus *bool `json:"chat_ip_eligibility_status"` -} - type versionPath struct { stable string canary string @@ -48,6 +37,7 @@ type versionPath struct { beta string } +// getVersionAndUserDataPath 根据当前平台返回实际存在的 Edge 用户数据目录。 func getVersionAndUserDataPath() (map[string]string, error) { var paths versionPath @@ -99,6 +89,7 @@ func getVersionAndUserDataPath() (map[string]string, error) { return versionDataPaths, nil } +// shutdownEdge 关闭当前正在运行的 Edge 进程，并返回被关闭进程的可执行文件路径。 func shutdownEdge() ([]string, error) { var terminatedEdges []string procs, err := process.Processes() @@ -112,23 +103,16 @@ func shutdownEdge() ([]string, error) { continue } - var isEdge bool - if runtime.GOOS == "darwin" { - isEdge = strings.HasPrefix(name, "Microsoft Edge") - } else { - isEdge = name == "msedge" - } - - if !isEdge { + if !isEdgeProcessName(runtime.GOOS, name) { continue } - // Check if process is still running + // 仅处理仍在运行的主进程，避免误杀同名辅助进程。 if running, _ := p.IsRunning(); !running { continue } - // Skip if parent has same name (avoid killing helper processes) + // 如果父进程同名，通常说明当前是子进程或 helper，直接跳过。 if ppid, _ := p.Ppid(); ppid > 0 { if parent, err := process.NewProcess(ppid); err == nil { if parentName, err := parent.Name(); err == nil && parentName == name { @@ -150,6 +134,7 @@ func shutdownEdge() ([]string, error) { return terminatedEdges, nil } +// getLastVersion 读取 Edge 用户目录中的 Last Version 文件。 func getLastVersion(userDataPath string) (string, error) { lastVersionFile := filepath.Join(userDataPath, "Last Version") if _, err := os.Stat(lastVersionFile); os.IsNotExist(err) { @@ -164,30 +149,21 @@ func getLastVersion(userDataPath string) (string, error) { return strings.TrimSpace(string(data)), nil } +// patchLocalState 仅修改 variations_country，同时保留 Local State 的其他字段。 func patchLocalState(userDataPath string) error { localStateFile := filepath.Join(userDataPath, "Local State") if _, err := os.Stat(localStateFile); os.IsNotExist(err) { return fmt.Errorf("failed to patch Local State. File not found: %s", localStateFile) } - data, err := os.ReadFile(localStateFile) + localState, err := readJSONObject(localStateFile) if err != nil { return err } - var localState LocalState - if err := json.Unmarshal(data, &localState); err != nil { - return err - } - - if localState.VariationsCountry != "US" { - localState.VariationsCountry = "US" - updatedData, err := json.MarshalIndent(localState, "", " ") - if err != nil { - return err - } - - if err := os.WriteFile(localStateFile, updatedData, 0644); err != nil { + if localState["variations_country"] != "US" { + localState["variations_country"] = "US" + if err := writeJSONObject(localStateFile, localState); err != nil { return err } fmt.Println("Succeeded in patching Local State") @@ -198,6 +174,7 @@ func patchLocalState(userDataPath string) error { return nil } +// patchPreferences 只处理 Default 和 Profile * 目录下的 Preferences 文件。 func patchPreferences(userDataPath string) error { entries, err := os.ReadDir(userDataPath) if err != nil { @@ -205,7 +182,7 @@ func patchPreferences(userDataPath string) error { } for _, entry := range entries { - if !entry.IsDir() && entry.Name() != "Default" && !strings.HasPrefix(entry.Name(), "Profile ") { + if !entry.IsDir() || !isTargetProfileDir(entry.Name()) { continue } @@ -214,29 +191,17 @@ func patchPreferences(userDataPath string) error { continue } - data, err := os.ReadFile(preferencesFile) + preferences, err := readJSONObject(preferencesFile) if err != nil { continue } - var preferences Preferences - if err := json.Unmarshal(data, &preferences); err != nil { - continue - } - - // Check if we need to patch - if preferences.Browser.ChatIPEligibilityStatus == nil || !*preferences.Browser.ChatIPEligibilityStatus { - // Set to true - trueVal := true - preferences.Browser.ChatIPEligibilityStatus = &trueVal + // browser 可能不存在，这里按需补齐，避免覆盖其他顶层字段。 + browser := ensureObject(preferences, "browser") - updatedData, err := json.MarshalIndent(preferences, "", " ") - if err != nil { - fmt.Printf("Failed to marshal preferences for %s: %v\n", entry.Name(), err) - continue - } - - if err := os.WriteFile(preferencesFile, updatedData, 0644); err != nil { + if browser["chat_ip_eligibility_status"] == nil || browser["chat_ip_eligibility_status"] == false { + browser["chat_ip_eligibility_status"] = true + if err := writeJSONObject(preferencesFile, preferences); err != nil { fmt.Printf("Failed to write preferences for %s: %v\n", entry.Name(), err) continue } @@ -249,27 +214,109 @@ func patchPreferences(userDataPath string) error { return nil } +// restartEdge 负责用已记录的可执行路径重新拉起 Edge。 func restartEdge(terminatedEdges []string) { - for _, edge := range terminatedEdges { - var cmd *string - switch runtime.GOOS { - case "windows": - cmdStr := edge + " --start-maximized" - cmd = &cmdStr - case "darwin": - cmdStr := "open -a 'Microsoft Edge' --args --start-maximized" - cmd = &cmdStr - case "linux": - cmdStr := edge + " --start-maximized" - cmd = &cmdStr + if err := restartEdges(terminatedEdges, runtime.GOOS, func(name string, args ...string) error { + cmd := exec.Command(name, args...) + cmd.Stderr = io.Discard + return cmd.Start() + }); err != nil { + fmt.Fprintf(os.Stderr, "Warning: Failed to restart Edge: %v\n", err) + } +} + +// isEdgeProcessName 兼容 msedge 与 msedge.exe，并保留 macOS 的前缀匹配行为。 +func isEdgeProcessName(goos string, name string) bool { + if goos == "darwin" { + return strings.HasPrefix(name, "Microsoft Edge") + } + + base := strings.TrimSuffix(strings.ToLower(name), strings.ToLower(filepath.Ext(name))) + return base == "msedge" +} + +// isTargetProfileDir 与 Python 版保持一致，只处理 Default 和 Profile *。 +func isTargetProfileDir(name string) bool { + return name == "Default" || strings.HasPrefix(name, "Profile ") +} + +// readJSONObject 读取 JSON 对象文件，便于按字段做最小修改。 +func readJSONObject(path string) (map[string]any, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("read JSON file %s: %w", path, err) + } + + var value map[string]any + if err := json.Unmarshal(data, &value); err != nil { + return nil, fmt.Errorf("unmarshal JSON file %s: %w", path, err) + } + + return value, nil +} + +// writeJSONObject 将完整对象写回文件，避免像结构体精简反序列化那样丢字段。 +func writeJSONObject(path string, value map[string]any) error { + data, err := json.Marshal(value) + if err != nil { + return fmt.Errorf("marshal JSON file %s: %w", path, err) + } + if err := os.WriteFile(path, data, 0644); err != nil { + return fmt.Errorf("write JSON file %s: %w", path, err) + } + return nil +} + +// ensureObject 保证指定键是对象；若不存在或类型不对，则创建一个空对象。 +func ensureObject(parent map[string]any, key string) map[string]any { + if value, ok := parent[key].(map[string]any); ok { + return value + } + + child := make(map[string]any) + parent[key] = child + return child +} + +// restartEdges 会先去重，再按平台构造命令逐个重启 Edge。 +func restartEdges(terminatedEdges []string, goos string, runner func(name string, args ...string) error) error { + for _, edge := range uniqueStrings(terminatedEdges) { + name, args, ok := buildRestartCommand(edge, goos) + if !ok { + continue } - if cmd != nil { - fmt.Printf("Starting: %s\n", edge) - // Note: In a real implementation, you'd use exec.Command here - // This is simplified for the example + fmt.Printf("Starting: %s\n", edge) + if err := runner(name, args...); err != nil { + return fmt.Errorf("restart edge %s: %w", edge, err) } } + + return nil +} + +// buildRestartCommand 将“可执行文件路径 + 参数”拆开，方便测试和注入 runner。 +func buildRestartCommand(edgePath, goos string) (string, []string, bool) { + switch goos { + case "windows", "linux", "darwin": + return edgePath, []string{"--start-maximized"}, true + default: + return "", nil, false + } +} + +// uniqueStrings 保留原顺序去重，避免同一个 Edge 可执行文件被重复启动。 +func uniqueStrings(values []string) []string { + seen := make(map[string]struct{}, len(values)) + result := make([]string, 0, len(values)) + for _, value := range values { + if _, ok := seen[value]; ok { + continue + } + seen[value] = struct{}{} + result = append(result, value) + } + return result } func main() { @@ -320,6 +367,9 @@ func main() { restartEdge(terminatedEdges) } + // 保持与原脚本一致：执行结束后等待用户确认，便于查看输出。 fmt.Print("\nPress Enter to continue...") - fmt.Scanln() + if _, err := fmt.Scanln(); err != nil { + fmt.Fprintf(os.Stderr, "Warning: failed to read input: %v\n", err) + } } diff --git a/main_test.go b/main_test.go new file mode 100644 index 0000000..23d3082 --- /dev/null +++ b/main_test.go @@ -0,0 +1,205 @@ +package main + +import ( + "encoding/json" + "os" + "path/filepath" + "reflect" + "testing" +) + +func TestPatchLocalStatePreservesExistingFields(t *testing.T) { + t.Parallel() + + tempDir := t.TempDir() + localStatePath := filepath.Join(tempDir, "Local State") + original := map[string]any{ + "variations_country": "CN", + "last_known_google_url": "https://example.com", + "nested": map[string]any{ + "enabled": true, + }, + } + writeJSONFile(t, localStatePath, original) + + if err := patchLocalState(tempDir); err != nil { + t.Fatalf("patchLocalState() error = %v", err) + } + + got := readJSONFile(t, localStatePath) + if got["variations_country"] != "US" { + t.Fatalf("variations_country = %v, want US", got["variations_country"]) + } + if got["last_known_google_url"] != original["last_known_google_url"] { + t.Fatalf("last_known_google_url = %v, want %v", got["last_known_google_url"], original["last_known_google_url"]) + } + if !reflect.DeepEqual(got["nested"], original["nested"]) { + t.Fatalf("nested = %#v, want %#v", got["nested"], original["nested"]) + } +} + +func TestPatchPreferencesOnlyUpdatesTargetProfilesAndPreservesFields(t *testing.T) { + t.Parallel() + + tempDir := t.TempDir() + defaultPath := filepath.Join(tempDir, "Default", "Preferences") + profilePath := filepath.Join(tempDir, "Profile 1", "Preferences") + systemProfilePath := filepath.Join(tempDir, "System Profile", "Preferences") + + originalDefault := map[string]any{ + "browser": map[string]any{ + "chat_ip_eligibility_status": false, + "show_home_button": true, + }, + "session": map[string]any{ + "restore_on_startup": 4, + }, + } + originalProfile := map[string]any{ + "browser": map[string]any{ + "show_home_button": false, + }, + "profile": map[string]any{ + "name": "Profile 1", + }, + } + originalSystemProfile := map[string]any{ + "browser": map[string]any{ + "chat_ip_eligibility_status": false, + }, + "system": true, + } + + writeJSONFile(t, defaultPath, originalDefault) + writeJSONFile(t, profilePath, originalProfile) + writeJSONFile(t, systemProfilePath, originalSystemProfile) + + if err := patchPreferences(tempDir); err != nil { + t.Fatalf("patchPreferences() error = %v", err) + } + + defaultGot := readJSONFile(t, defaultPath) + assertChatEligibility(t, defaultGot, true) + assertNestedField(t, defaultGot, "browser", "show_home_button", true) + assertNestedField(t, defaultGot, "session", "restore_on_startup", float64(4)) + + profileGot := readJSONFile(t, profilePath) + assertChatEligibility(t, profileGot, true) + assertNestedField(t, profileGot, "browser", "show_home_button", false) + assertNestedField(t, profileGot, "profile", "name", "Profile 1") + + systemProfileGot := readJSONFile(t, systemProfilePath) + if !reflect.DeepEqual(systemProfileGot, originalSystemProfile) { + t.Fatalf("System Profile changed = %#v, want %#v", systemProfileGot, originalSystemProfile) + } +} + +func TestIsEdgeProcessName(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + goos string + process string + want bool + }{ + {name: "windows exe", goos: "windows", process: "msedge.exe", want: true}, + {name: "windows bare", goos: "windows", process: "msedge", want: true}, + {name: "linux bare", goos: "linux", process: "msedge", want: true}, + {name: "linux exe", goos: "linux", process: "msedge.exe", want: true}, + {name: "darwin app", goos: "darwin", process: "Microsoft Edge Helper", want: true}, + {name: "other process", goos: "linux", process: "chrome", want: false}, + } + + for _, tt := range tests { + tt := tt + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + if got := isEdgeProcessName(tt.goos, tt.process); got != tt.want { + t.Fatalf("isEdgeProcessName(%q, %q) = %v, want %v", tt.goos, tt.process, got, tt.want) + } + }) + } +} + +func TestRestartEdgesRunsUniqueCommands(t *testing.T) { + t.Parallel() + + var calls []string + runner := func(name string, args ...string) error { + call := name + if len(args) > 0 { + call += " " + args[0] + } + calls = append(calls, call) + return nil + } + + if err := restartEdges([]string{"/opt/msedge", "/opt/msedge", "/opt/msedge-beta"}, "linux", runner); err != nil { + t.Fatalf("restartEdges() error = %v", err) + } + + want := []string{ + "/opt/msedge --start-maximized", + "/opt/msedge-beta --start-maximized", + } + if !reflect.DeepEqual(calls, want) { + t.Fatalf("runner calls = %#v, want %#v", calls, want) + } +} + +func writeJSONFile(t *testing.T, path string, value map[string]any) { + t.Helper() + + if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil { + t.Fatalf("MkdirAll(%q) error = %v", path, err) + } + + data, err := json.Marshal(value) + if err != nil { + t.Fatalf("json.Marshal() error = %v", err) + } + if err := os.WriteFile(path, data, 0644); err != nil { + t.Fatalf("WriteFile(%q) error = %v", path, err) + } +} + +func readJSONFile(t *testing.T, path string) map[string]any { + t.Helper() + + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("ReadFile(%q) error = %v", path, err) + } + + var value map[string]any + if err := json.Unmarshal(data, &value); err != nil { + t.Fatalf("json.Unmarshal(%q) error = %v", path, err) + } + + return value +} + +func assertChatEligibility(t *testing.T, value map[string]any, want bool) { + t.Helper() + + browser, ok := value["browser"].(map[string]any) + if !ok { + t.Fatalf("browser = %#v, want map", value["browser"]) + } + if got, ok := browser["chat_ip_eligibility_status"].(bool); !ok || got != want { + t.Fatalf("chat_ip_eligibility_status = %#v, want %v", browser["chat_ip_eligibility_status"], want) + } +} + +func assertNestedField(t *testing.T, value map[string]any, firstKey, secondKey string, want any) { + t.Helper() + + nested, ok := value[firstKey].(map[string]any) + if !ok { + t.Fatalf("%s = %#v, want map", firstKey, value[firstKey]) + } + if got := nested[secondKey]; !reflect.DeepEqual(got, want) { + t.Fatalf("%s.%s = %#v, want %#v", firstKey, secondKey, got, want) + } +}