State + Exports: Make cleanup optional when entering disconnect

4 files changed, 21 insertions, 19 deletions

diff --git a/exports/exports.go b/exports/exports.go
index 3ebd7f0..a6dbaf6 100644
--- a/exports/exports.go
+++ b/exports/exports.go
@@ -244,13 +244,13 @@ func SetSearchServer(name *C.char) *C.char {
 }
 
 //export SetDisconnected
-func SetDisconnected(name *C.char) *C.char {
+func SetDisconnected(name *C.char, cleanup C.int) *C.char {
 	nameStr := C.GoString(name)
 	state, stateErr := GetVPNState(nameStr)
 	if stateErr != nil {
 		return C.CString(ErrorToString(stateErr))
 	}
-	setDisconnectedErr := state.SetDisconnected()
+	setDisconnectedErr := state.SetDisconnected(int(cleanup) == 1)
 	return C.CString(ErrorToString(setDisconnectedErr))
 }
 
diff --git a/state.go b/state.go
index fdfdd45..08dc189 100644
--- a/state.go
+++ b/state.go
@@ -390,7 +390,7 @@ func (state *VPNState) SetDisconnecting() error {
 	return nil
 }
 
-func (state *VPNState) SetDisconnected() error {
+func (state *VPNState) SetDisconnected(cleanup bool) error {
 	errorMessage := "failed to set disconnected"
 	if state.InFSMState(fsm.HAS_CONFIG) {
 		// already disconnected, show no error
@@ -400,22 +400,24 @@ func (state *VPNState) SetDisconnected() error {
 		return &types.WrappedErrorMessage{Message: errorMessage, Err: fsm.WrongStateTransitionError{Got: state.FSM.Current, Want: fsm.HAS_CONFIG}.CustomError()}
 	}
 
-	// Do the /disconnect API call and go to disconnected after...
-	currentServer, currentServerErr := state.Servers.GetCurrentServer()
-	if currentServerErr != nil {
-		return &types.WrappedErrorMessage{Message: errorMessage, Err: currentServerErr}
-	}
+	if cleanup {
+		// Do the /disconnect API call and go to disconnected after...
+		currentServer, currentServerErr := state.Servers.GetCurrentServer()
+		if currentServerErr != nil {
+			return &types.WrappedErrorMessage{Message: errorMessage, Err: currentServerErr}
+		}
 
-	oauthStructure := currentServer.GetOAuth()
+		oauthStructure := currentServer.GetOAuth()
 
-	// Make sure the FSM is initialized
-	oauthStructure.FSM = &state.FSM
-	base, baseErr := currentServer.GetBase()
-	if baseErr != nil {
-		return &types.WrappedErrorMessage{Message: errorMessage, Err: baseErr}
+		// Make sure the FSM is initialized
+		oauthStructure.FSM = &state.FSM
+		base, baseErr := currentServer.GetBase()
+		if baseErr != nil {
+			return &types.WrappedErrorMessage{Message: errorMessage, Err: baseErr}
+		}
+		base.FSM = &state.FSM
+		server.Disconnect(currentServer)
 	}
-	base.FSM = &state.FSM
-	server.Disconnect(currentServer)
 
 	state.FSM.GoTransitionWithData(fsm.HAS_CONFIG, state.getServerInfoData(), false)
 
diff --git a/wrappers/python/src/__init__.py b/wrappers/python/src/__init__.py
index 761adf5..74a529d 100644
--- a/wrappers/python/src/__init__.py
+++ b/wrappers/python/src/__init__.py
@@ -85,7 +85,7 @@ lib.SetSecureLocation.argtypes, lib.SetSecureLocation.restype = [
 lib.SetConnected.argtypes, lib.SetConnected.restype = [c_char_p], c_void_p
 lib.SetDisconnecting.argtypes, lib.SetDisconnecting.restype = [c_char_p], c_void_p
 lib.SetConnecting.argtypes, lib.SetConnecting.restype = [c_char_p], c_void_p
-lib.SetDisconnected.argtypes, lib.SetDisconnected.restype = [c_char_p], c_void_p
+lib.SetDisconnected.argtypes, lib.SetDisconnected.restype = [c_char_p, c_int], c_void_p
 lib.SetSearchServer.argtypes, lib.SetSearchServer.restype = [c_char_p], c_void_p
 lib.ShouldRenewButton.argtypes, lib.ShouldRenewButton.restype = [], int
 lib.RenewSession.argtypes, lib.RenewSession.restype = [c_char_p], c_void_p
diff --git a/wrappers/python/src/main.py b/wrappers/python/src/main.py
index 574d0e2..e70d431 100644
--- a/wrappers/python/src/main.py
+++ b/wrappers/python/src/main.py
@@ -157,8 +157,8 @@ class EduVPN(object):
         if connecting_err:
             raise Exception(connecting_err)
 
-    def set_disconnected(self) -> None:
-        disconnect_err = self.go_function(lib.SetDisconnected)
+    def set_disconnected(self, cleanup=True) -> None:
+        disconnect_err = self.go_function(lib.SetDisconnected, cleanup)
 
         if disconnect_err:
             raise Exception(disconnect_err)