From c05b73f8ecd9e4661f2cb86c8a9681ce6ff3f6d3 Mon Sep 17 00:00:00 2001
From: Vladislav Yarmak <vladislav-ex-src@vm-0.com>
Date: Sun, 8 Oct 2023 13:46:38 +0300
Subject: [PATCH] add CLI option to force cold init

---
 main.go | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index 25d48bf..67d3919 100644
--- a/main.go
+++ b/main.go
@@ -62,6 +62,7 @@ type CLIArgs struct {
 	password         string
 	tfacode          string
 	fakeSNI          string
+	forceColdInit    bool
 }
 
 func parse_args() CLIArgs {
@@ -90,6 +91,7 @@ func parse_args() CLIArgs {
 	flag.StringVar(&args.password, "password", "", "password for login")
 	flag.StringVar(&args.tfacode, "2fa", "", "2FA code for login")
 	flag.StringVar(&args.fakeSNI, "fake-sni", "com", "fake SNI to use to contact windscribe servers")
+	flag.BoolVar(&args.forceColdInit, "force-cold-init", false, "force cold init")
 	flag.Parse()
 	if args.listLocations && args.listProxies {
 		arg_fail("list-locations and list-proxies flags are mutually exclusive")
@@ -192,9 +194,14 @@ func run() int {
 	wndc.Mux.Unlock()
 
 	// Try ressurect state
-	state, err := loadState(args.stateFile)
+	state, err := maybeLoadState(args.forceColdInit, args.stateFile)
 	if err != nil {
-		mainLogger.Warning("Failed to load client state: %v. It is OK for a first run. Performing cold init...", err)
+		switch err {
+		case errColdInitForced:
+			mainLogger.Info("Cold init forced.")
+		default:
+			mainLogger.Warning("Failed to load client state: %v. It is OK for a first run. Performing cold init...", err)
+		}
 		err = coldInit(wndc, args.username, args.password, args.tfacode, args.timeout)
 		if err != nil {
 			mainLogger.Critical("Cold init failed: %v", err)
@@ -345,6 +352,14 @@ func pickServer(serverList wndclient.ServerList, location string) string {
 	return candidates[rnd.Intn(len(candidates))]
 }
 
+var errColdInitForced = errors.New("cold init forced!")
+func maybeLoadState(forceColdInit bool, filename string) (*wndclient.WndClientState, error) {
+	if forceColdInit {
+		return nil, errColdInitForced
+	}
+	return loadState(filename)
+}
+
 func loadState(filename string) (*wndclient.WndClientState, error) {
 	file, err := os.Open(filename)
 	if err != nil {