diff --git a/cmd/root.go b/cmd/root.go index 31af1cf..1eac1ac 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -2,8 +2,11 @@ package cmd import ( "fmt" + "log" "os" + "path/filepath" + "github.com/nicksherron/bashhub-server/internal" "github.com/spf13/cobra" ) @@ -11,21 +14,12 @@ var cfgFile string // rootCmd represents the base command when called without any subcommands var rootCmd = &cobra.Command{ - Use: "generated code example", - Short: "A brief description of your application", - Long: `A longer description that spans multiple lines and likely contains -examples and usage of using your application. For example: - -Cobra is a CLI library for Go that empowers applications. -This application is a tool to generate the needed files -to quickly create a Cobra application.`, - // Uncomment the following line if your bare application - // has an action associated with it: - // Run: func(cmd *cobra.Command, args []string) { }, + Run: func(cmd *cobra.Command, args []string) { + cmd.Flags().Parse(args) + internal.Run() + }, } -// Execute adds all child commands to the root command and sets flags appropriately. -// This is called by main.main(). It only needs to happen once to the rootCmd. func Execute() { if err := rootCmd.Execute(); err != nil { fmt.Println(err) @@ -35,9 +29,26 @@ func Execute() { func init() { cobra.OnInitialize() - - // Cobra also supports local flags, which will only run - // when this action is called directly. - // rootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle") + rootCmd.PersistentFlags().StringVar(&internal.DbPath, "db", dbPath(), "DB location (sqlite or postgres)") } +func dbPath() string { + dbFile := "data.db" + f := filepath.Join(appDir(), dbFile) + return f +} + +func appDir() string { + cfgDir, err := os.UserConfigDir() + if err != nil { + log.Fatal(err) + } + + ch := filepath.Join(cfgDir, ".bashhub-server") + err = os.MkdirAll(ch, 0755) + if err != nil { + log.Fatal(err) + } + + return ch +} diff --git a/go.mod b/go.mod index 0047406..f87ca00 100644 --- a/go.mod +++ b/go.mod @@ -1,7 +1,15 @@ module github.com/nicksherron/bashhub-server require ( - + github.com/dgrijalva/jwt-go v3.2.0+incompatible + github.com/gin-gonic/contrib v0.0.0-20191209060500-d6e26eeaa607 + github.com/gin-gonic/gin v1.5.0 + github.com/inconshreveable/mousetrap v1.0.0 // indirect + github.com/jinzhu/gorm v1.9.12 + github.com/lacion/cookiecutter_golang_example v0.0.0-20191209145422-f4f6c7d38761 + github.com/lib/pq v1.3.0 github.com/spf13/cobra v0.0.3 - + github.com/spf13/pflag v1.0.5 // indirect ) + +go 1.13 diff --git a/internal/db.go b/internal/db.go new file mode 100644 index 0000000..52507c4 --- /dev/null +++ b/internal/db.go @@ -0,0 +1,179 @@ +package internal + +import ( + "database/sql" + "fmt" + "log" + "strings" + "time" + + "github.com/jinzhu/gorm" + _ "github.com/jinzhu/gorm/dialects/postgres" + _ "github.com/jinzhu/gorm/dialects/sqlite" + _ "github.com/lib/pq" +) + +var ( + DB *sql.DB + DbPath string + connectionLimit int +) + +// DbInit initializes our db. +func DbInit() { + // GormDB contains DB connection state + var gormdb *gorm.DB + + var err error + if strings.HasPrefix(DbPath, "postgres://") { + // + DB, err = sql.Open("postgres", DbPath) + if err != nil { + log.Fatal(err) + } + + gormdb, err = gorm.Open("postgres", DbPath) + if err != nil { + log.Fatal(err) + } + connectionLimit = 50 + } else { + DbPath = fmt.Sprintf("file:%v?cache=shared&mode=rwc", DbPath) + DB, err = sql.Open("sqlite3", DbPath) + if err != nil { + log.Fatal(err) + } + gormdb, err = gorm.Open("sqlite3", DbPath) + if err != nil { + log.Fatal(err) + } + DB.Exec("PRAGMA journal_mode=WAL;") + connectionLimit = 1 + + } + DB.SetMaxOpenConns(connectionLimit) + gormdb.AutoMigrate(&User{}) + gormdb.AutoMigrate(&Command{}) + gormdb.AutoMigrate(&System{}) + gormdb.Model(&User{}).AddIndex("idx_user", "username") + gormdb.Model(&User{}).AddIndex("idx_token", "token") + gormdb.Model(&System{}).AddIndex("idx_mac", "mac") + + // just need gorm for migration. + gormdb.Close() +} + +func (user User) userExists() bool { + var exists bool + err := DB.QueryRow("SELECT exists (select id from users where username = $1 and password = $2)", + user.Username, user.Password).Scan(&exists) + if err != nil && err != sql.ErrNoRows { + log.Fatalf("error checking if row exists %v", err) + } + return exists +} + +func (user User) userCreate() int64 { + res, err := DB.Exec(`INSERT into users("registration_code", "username","password","email") + VALUES ($1,$2,$3,$4) ON CONFLICT(username) do nothing`, user.RegistrationCode, + user.Username, user.Password, user.Email) + if err != nil { + log.Fatal(err) + } + inserted, err := res.RowsAffected() + if err != nil { + log.Fatal(err) + } + return inserted +} + +func (user User) updateToken() { + _, err := DB.Exec(`UPDATE users set "token" = $1 where "username" = $2 `, user.Token, user.Username) + if err != nil { + log.Fatal(err) + } +} + +func (user User) tokenExists() bool { + var exists bool + err := DB.QueryRow("SELECT exists (select id from users where token = $1)", + user.Token).Scan(&exists) + if err != nil && err != sql.ErrNoRows { + log.Fatalf("error checking if row exists %v", err) + } + return exists +} + +func (cmd Command) commandInsert() int64 { + res, err := DB.Exec(`INSERT into commands("uuid", "command", "created", "user_id") + VALUES ($1,$2,$3,(select "id" from users where "token" = $4))`, + cmd.Uuid, cmd.Command, cmd.Created, cmd.Token) + if err != nil { + log.Fatal(err) + } + inserted, err := res.RowsAffected() + if err != nil { + log.Fatal(err) + } + return inserted +} + +func (cmd Command) commandGet() []Query { + var results []Query + var rows *sql.Rows + var err error + if cmd.Unique { + rows, err = DB.Query(`SELECT "command", "uuid", "created" from commands + where "user_id" in (select "id" from users where "token" = $1) + group by command order by created desc limit $2`, + cmd.Token, cmd.Limit) + } else { + rows, err = DB.Query(`SELECT "command", "uuid", "created" from commands + where "user_id" in (select "id" from users where "token" = $1) order by created desc limit $2`, + cmd.Token, cmd.Limit) + } + + if err != nil { + log.Println(err) + } + defer rows.Close() + for rows.Next() { + var result Query + err = rows.Scan(&result.Command, &result.Uuid, &result.Created) + if err != nil { + log.Println(err) + } + results = append(results, result) + } + + return results + +} + +func (sys System) systemInsert() int64 { + + t := time.Now().Unix() + res, err := DB.Exec(`INSERT into systems ("name", "mac", "user_id", "hostname", "client_version", "created", "updated") + VALUES ($1, $2, (select "id" from users where "token" = $3), $4, $5, $6, $7)`, + sys.Name, sys.Mac, sys.Token, sys.Hostname, sys.ClientVersion, t, t) + if err != nil { + log.Fatal(err) + } + inserted, err := res.RowsAffected() + if err != nil { + log.Fatal(err) + } + return inserted +} + +func (sys System) systemGet() (SystemQuery, error) { + var row SystemQuery + err := DB.QueryRow(`SELECT "name", "mac", "user_id", "hostname", "client_version", + "id", "created", "updated" from systems where mac = $1`, + sys.Mac).Scan(&row) + if err != nil { + return SystemQuery{}, err + } + return row, nil + +} diff --git a/internal/server.go b/internal/server.go new file mode 100644 index 0000000..c80ec37 --- /dev/null +++ b/internal/server.go @@ -0,0 +1,207 @@ +package internal + +import ( + "crypto/sha256" + "fmt" + "log" + "net/http" + "strconv" + "strings" + "time" + + jwt_lib "github.com/dgrijalva/jwt-go" + "github.com/gin-gonic/gin" +) + +type User struct { + ID uint `gorm:"primary_key"` + Username string `form:"Username" json:"Username" xml:"Username" gorm:"type:varchar(200);unique_index"` + Email string `form:"email" json:"email" xml:"email"` + Password string `form:"password" json:"password" xml:"password"` + Mac *string `form:"mac" json:"mac" xml:"mac"` + RegistrationCode *string `form:"registrationCode" json:"registrationCode" xml:"registrationCode"` + Token string +} + +type Query struct { + Uuid string `form:"uuid" json:"uuid" xml:"uuid"` + Command string `form:"command" json:"command" xml:"command"` + Created int64 `form:"created" json:"created" xml:"created"` +} +type SystemQuery struct { + ID uint `form:"id" json:"id" xml:"id" gorm:"primary_key"` + Created int64 + Updated int64 + Mac string `form:"mac" json:"mac" xml:"mac"` + Hostname *string `form:"hostname" json:"hostname" xml:"hostname"` + Name *string `form:"name" json:"name" xml:"name"` + ClientVersion *string `form:"clientVersion" json:"clientVersion" xml:"clientVersion"` +} + +type Command struct { + ProcessId int `form:"processId" json:"processId" xml:"processId"` + ProcessStartTime int64 `form:"processStartTime" json:"processStartTime" xml:"processStartTime"` + Uuid string `form:"uuid" json:"uuid" xml:"uuid"` + Command string `form:"command" json:"command" xml:"command"` + Created int64 `form:"created" json:"created" xml:"created"` + Path string `form:"path" json:"path" xml:"path"` + ExitStatus int `form:"exitStatus" json:"exitStatus" xml:"exitStatus"` + User User `gorm:"association_foreignkey:ID"` + UserId uint + Token string `gorm:"-"` + Limit int `gorm:"-"` + Unique bool `gorm:"-"` +} + +// {"mac": "83779604164095", "hostname": "yay.local", "name": "yay.local", "clientVersion": "1.2.0"} +//{"name":"Home","mac":"83779604164095","userId":"5b5d53b6e4b02a6c4914bec8","hostname":"yay.local","clientVersion":"1.2.0","id":"5b5d53c8e4b02a6c4914bec9","created":1532842952382,"updated":1581032237766} +type System struct { + ID uint `form:"id" json:"id" xml:"id" gorm:"primary_key"` + Created int64 + Updated int64 + Mac string `form:"mac" json:"mac" xml:"mac"` + Hostname *string `form:"hostname" json:"hostname" xml:"hostname"` + Name *string `form:"name" json:"name" xml:"name"` + ClientVersion *string `form:"clientVersion" json:"clientVersion" xml:"clientVersion"` + User User `gorm:"association_foreignkey:ID"` + UserId uint `form:"userId" json:"userId" xml:"userId"` + Token string `gorm:"-"` +} + +func auth() gin.HandlerFunc { + return func(c *gin.Context) { + var user User + err := func() error { + user.Token = strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ") + if user.tokenExists() { + return nil + } else { + return fmt.Errorf("token doesn't exist") + } + }() + if err != nil { + c.AbortWithError(401, err) + } + } +} + +func Run() { + + DbInit() + + r := gin.Default() + + r.GET("/ping", func(c *gin.Context) { + c.JSON(200, gin.H{ + "message": "pong", + }) + }) + + r.POST("/api/v1/login", func(c *gin.Context) { + var user User + if err := c.ShouldBindJSON(&user); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + user.Password = fmt.Sprintf("%v", sha256.Sum256([]byte(user.Password))) + + if !user.userExists() { + c.String(401, "Bad credentials") + return + } + + token := jwt_lib.New(jwt_lib.GetSigningMethod("HS256")) + + token.Claims = jwt_lib.MapClaims{ + "Id": user.Username, + "exp": time.Now().Add(time.Hour * 20000).Unix(), + } + // Sign and get the complete encoded token as a string + tokenString, err := token.SignedString([]byte(user.Password)) + if err != nil { + c.JSON(500, gin.H{"message": "Could not generate token"}) + } + user.Token = tokenString + user.updateToken() + + c.JSON(200, gin.H{"accessToken": tokenString}) + + }) + + r.POST("/api/v1/user", func(c *gin.Context) { + var user User + if err := c.ShouldBindJSON(&user); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if user.Email == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "email required"}) + return + } + + user.Password = fmt.Sprintf("%v", sha256.Sum256([]byte(user.Password))) + + user.userCreate() + + }) + r.Use(auth()) + r.GET("/api/v1/command/search", func(c *gin.Context) { + var command Command + command.Token = strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ") + command.Limit = 100 + if c.Query("limit") != "" { + if num, err := strconv.Atoi(c.Query("limit")); err != nil { + command.Limit = num + } + } + if c.Query("unique") == "true" { + command.Unique = true + }else { + command.Unique = false + } + result := command.commandGet() + c.IndentedJSON(http.StatusOK, result) + + }) + + r.POST("/api/v1/command", func(c *gin.Context) { + var command Command + if err := c.ShouldBindJSON(&command); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + command.Token = strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ") + command.commandInsert() + }) + + r.POST("/api/v1/system", func(c *gin.Context) { + var system System + err := c.Bind(&system) + if err != nil { + log.Fatal(err) + } + system.Token = strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ") + system.systemInsert() + c.AbortWithStatus(201) + }) + + r.GET("/api/v1/system", func(c *gin.Context) { + var system System + system.Token = strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ") + system.Mac = c.Query("mac") + if system.Mac == "" { + c.AbortWithStatus(http.StatusBadRequest) + return + + } + result, err := system.systemGet() + if err != nil { + c.AbortWithStatus(404) + return + } + c.IndentedJSON(http.StatusOK, result) + + }) + r.Run() + +} diff --git a/main.go b/main.go index 043ca04..30f5d62 100644 --- a/main.go +++ b/main.go @@ -1,13 +1,11 @@ package main import ( - "github.com/nicksherron/bashhub-server/cmd" ) func main() { - - cmd.Execute() - + cmd.Execute() + }