Server + State + Types: Return marshalled JSON for display information

3 files changed, 105 insertions, 17 deletions

diff --git a/internal/server/server.go b/internal/server/server.go
index f88c521..a34e8ec 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -15,11 +15,13 @@ import (
 // The base type for servers
 type ServerBase struct {
 	URL         string            `json:"base_url"`
+	DisplayName map[string]string `json:"display_name"`
+	SupportContact []string `json:"support_contact"`
 	Endpoints   ServerEndpoints   `json:"endpoints"`
 	Profiles    ServerProfileInfo `json:"profiles"`
 	ProfilesRaw string            `json:"profiles_raw"`
-	StartTime   int64             `json:"start-time"`
-	EndTime     int64             `json:"end-time"`
+	StartTime   int64             `json:"start_time"`
+	EndTime     int64             `json:"expire_time"`
 	Logger      *log.FileLogger   `json:"-"`
 	FSM         *fsm.FSM          `json:"-"`
 }
@@ -36,7 +38,7 @@ type InstituteAccessServer struct {
 // A secure internet server which has its own OAuth tokens
 // It specifies the current location url it is connected to
 type SecureInternetHomeServer struct {
-	DisplayName string      `json:"display_name"`
+	DisplayName map[string]string      `json:"display_name"`
 	OAuth       oauth.OAuth `json:"oauth"`
 
 	// The home server has a list of info for each configured server location
@@ -56,6 +58,9 @@ type InstituteServers struct {
 func (servers *Servers) GetCurrentServer() (Server, error) {
 	errorMessage := "failed getting current server"
 	if servers.IsSecureInternet {
+		if !servers.HasSecureLocation() {
+			return nil, &types.WrappedErrorMessage{Message: errorMessage, Err: &ServerGetCurrentNotFoundError{}}
+		}
 		return &servers.SecureInternetHomeServer, nil
 	}
 	currentInstitute := servers.InstituteServers.CurrentURL
@@ -137,9 +142,11 @@ func (server *SecureInternetHomeServer) GetBase() (*ServerBase, error) {
 	return base, nil
 }
 
-func (institute *InstituteAccessServer) init(url string, fsm *fsm.FSM, logger *log.FileLogger) error {
+func (institute *InstituteAccessServer) init(url string, displayName map[string]string, supportContact []string, fsm *fsm.FSM, logger *log.FileLogger) error {
 	errorMessage := fmt.Sprintf("failed initializing institute server %s", url)
 	institute.Base.URL = url
+	institute.Base.DisplayName = displayName
+	institute.Base.SupportContact = supportContact
 	institute.Base.FSM = fsm
 	institute.Base.Logger = logger
 	endpoints, endpointsErr := APIGetEndpoints(url)
@@ -169,6 +176,8 @@ func (secure *SecureInternetHomeServer) addLocation(locationServer *types.Discov
 		// Create the base to be added to the map
 		base = &ServerBase{}
 		base.URL = locationServer.BaseURL
+		base.DisplayName = secure.DisplayName
+		base.SupportContact = locationServer.SupportContact
 		endpoints, endpointsErr := APIGetEndpoints(locationServer.BaseURL)
 		if endpointsErr != nil {
 			return nil, &types.WrappedErrorMessage{Message: errorMessage, Err: endpointsErr}
@@ -194,17 +203,19 @@ func (secure *SecureInternetHomeServer) init(homeOrg *types.DiscoveryOrganizatio
 		*secure = *&SecureInternetHomeServer{}
 	}
 
+	// Make sure to set the organization ID
+	secure.HomeOrganizationID = homeOrg.OrgId
+	secure.DisplayName = homeOrg.DisplayName
+
+	// Make sure to set the authorization URL template
+	secure.AuthorizationTemplate = homeLocation.AuthenticationURLTemplate
+
 	base, baseErr := secure.addLocation(homeLocation, fsm, logger)
 
 	if baseErr != nil {
 		return &types.WrappedErrorMessage{Message: errorMessage, Err: baseErr}
 	}
 
-	// Make sure to set the organization ID
-	secure.HomeOrganizationID = homeOrg.OrgId
-
-	// Make sure to set the authorization URL template
-	secure.AuthorizationTemplate = homeLocation.AuthenticationURLTemplate
 	// Make sure oauth contains our endpoints
 	secure.OAuth.Init(base.Endpoints.API.V3.Authorization, base.Endpoints.API.V3.Token, fsm, logger)
 	return nil
@@ -290,7 +301,7 @@ func (servers *Servers) AddInstituteAccess(instituteServer *types.DiscoveryServe
 
 	// Set the current server
 	instituteServers.CurrentURL = url
-	instituteInitErr := institute.init(url, fsm, logger)
+	instituteInitErr := institute.init(url, instituteServer.DisplayName, instituteServer.SupportContact, fsm, logger)
 	if instituteInitErr != nil {
 		return nil, &types.WrappedErrorMessage{Message: errorMessage, Err: instituteInitErr}
 	}
@@ -511,6 +522,49 @@ func askForProfileID(server Server) error {
 	return nil
 }
 
+type ServerInfoScreen struct {
+	DisplayName map[string]string `json:"display_name"`
+	CountryCode string `json:"country_code,omitempty"`
+	SupportContact []string `json:"support_contact"`
+	ProfilesRaw string `json:"profiles"`
+	ExpireTime int64 `json:"expire_time"`
+}
+
+func (servers *Servers) GetCurrentServerInfoJSON() (string, error) {
+	errorMessage := "failed getting JSON for server"
+
+
+	currentServer, currentServerErr := servers.GetCurrentServer()
+	if currentServerErr != nil {
+		return "{}", &types.WrappedErrorMessage{Message: errorMessage, Err: currentServerErr}
+	}
+
+	serverInfoScreen := &ServerInfoScreen{}
+
+	base, baseErr := currentServer.GetBase()
+
+	if baseErr != nil {
+		return "{}", &types.WrappedErrorMessage{Message: errorMessage, Err: baseErr}
+	}
+
+	serverInfoScreen.DisplayName = base.DisplayName
+	serverInfoScreen.SupportContact = base.SupportContact
+	serverInfoScreen.ProfilesRaw = base.ProfilesRaw
+	serverInfoScreen.ExpireTime = base.EndTime
+
+	if servers.IsSecureInternet {
+		serverInfoScreen.CountryCode = servers.SecureInternetHomeServer.CurrentLocation
+	}
+
+	bytes, bytesErr := json.Marshal(serverInfoScreen)
+
+	if bytesErr != nil {
+		return "{}", &types.WrappedErrorMessage{Message: errorMessage, Err: bytesErr}
+	}
+
+	return string(bytes), nil
+}
+
 func GetConfig(server Server, forceTCP bool) (string, string, error) {
 	errorMessage := "failed getting an OpenVPN/WireGuard configuration"
 	base, baseErr := server.GetBase()
diff --git a/internal/types/server.go b/internal/types/server.go
index 31a7bcc..09bfc92 100644
--- a/internal/types/server.go
+++ b/internal/types/server.go
@@ -1,5 +1,9 @@
 package types
 
+import (
+	"encoding/json"
+)
+
 // Shared server types
 
 // Structs that define the json format for
@@ -12,9 +16,7 @@ type DiscoveryOrganizations struct {
 }
 
 type DiscoveryOrganization struct {
-	DisplayName struct {
-		En string `json:"en"`
-	} `json:"display_name"`
+	DisplayName map[string]string `json:"display_name"`
 	OrgId              string `json:"org_id"`
 	SecureInternetHome string `json:"secure_internet_home"`
 	KeywordList        struct {
@@ -31,10 +33,37 @@ type DiscoveryServers struct {
 	RawString string            `json:"-"`
 }
 
+type DNMapOrString map[string]string
+
+// The display name can either be a map or a string in the server list
+// Unmarshal it by first trying a string and then the map
+func (DN *DNMapOrString) UnmarshalJSON(data []byte) error {
+	var displayNameString string
+
+	err := json.Unmarshal(data, &displayNameString)
+
+	if err == nil {
+		*DN = map[string]string{"en": displayNameString}
+		return nil
+	}
+
+	var resultingMap map[string]string
+
+	err = json.Unmarshal(data, &resultingMap)
+
+	if err == nil {
+		*DN = resultingMap
+		return nil
+	}
+	return err
+}
+
+
 type DiscoveryServer struct {
 	AuthenticationURLTemplate string   `json:"authentication_url_template"`
 	BaseURL                   string   `json:"base_url"`
 	CountryCode               string   `json:"country_code"`
+	DisplayName               DNMapOrString   `json:"display_name,omitempty"`
 	PublicKeyList             []string `json:"public_key_list"`
 	Type                      string   `json:"server_type"`
 	SupportContact            []string `json:"support_contact"`
diff --git a/state.go b/state.go
index b8fa046..e1b26ed 100644
--- a/state.go
+++ b/state.go
@@ -134,7 +134,7 @@ func (state *VPNState) getConfig(chosenServer server.Server, forceTCP bool) (str
 		state.GoBack()
 		return "", "", &types.WrappedErrorMessage{Message: errorMessage, Err: configErr}
 	} else {
-		state.FSM.GoTransition(fsm.HAS_CONFIG)
+		state.FSM.GoTransitionWithData(fsm.HAS_CONFIG, state.getServerInfoData(), false)
 	}
 
 	return config, configType, nil
@@ -230,7 +230,7 @@ func (state *VPNState) addInstituteServer(url string) (server.Server, error) {
 func (state *VPNState) addCustomServer(url string) (server.Server, error) {
 	errorMessage := fmt.Sprintf("failed adding Custom server with url %s", url)
 
-	instituteServer := &types.DiscoveryServer{BaseURL: url, CountryCode: "NL", Type: "custom", SupportContact: []string{"custom"}}
+	instituteServer := &types.DiscoveryServer{BaseURL: url, DisplayName: map[string]string{"en": url}, Type: "custom_server"}
 
 	// A custom server is just an institute access server under the hood
 	server, serverErr := state.Servers.AddInstituteAccess(instituteServer, &state.FSM, &state.Logger)
@@ -326,6 +326,11 @@ func (state *VPNState) SetSearchServer() error {
 	return nil
 }
 
+func (state *VPNState) getServerInfoData() string {
+	jsonString, _ := state.Servers.GetCurrentServerInfoJSON()
+	return jsonString
+}
+
 func (state *VPNState) SetConnected() error {
 	if state.FSM.InState(fsm.CONNECTED) {
 		// already connected, show no error
@@ -335,7 +340,7 @@ func (state *VPNState) SetConnected() error {
 		return &types.WrappedErrorMessage{Message: "failed to set connected", Err: fsm.WrongStateTransitionError{Got: state.FSM.Current, Want: fsm.CONNECTED}.CustomError()}
 	}
 
-	state.FSM.GoTransition(fsm.CONNECTED)
+	state.FSM.GoTransitionWithData(fsm.CONNECTED, state.getServerInfoData(), false)
 	return nil
 }
 
@@ -361,7 +366,7 @@ func (state *VPNState) SetDisconnected() error {
 		return &types.WrappedErrorMessage{Message: "failed to set disconnected", Err: fsm.WrongStateTransitionError{Got: state.FSM.Current, Want: fsm.HAS_CONFIG}.CustomError()}
 	}
 
-	state.FSM.GoTransition(fsm.HAS_CONFIG)
+	state.FSM.GoTransitionWithData(fsm.HAS_CONFIG, state.getServerInfoData(), false)
 	return nil
 }