package main import ( "encoding/json" "fmt" "io" "log" "net" "os" "strings" "gopkg.in/ini.v1" "gorm.io/driver/postgres" "gorm.io/driver/sqlite" "gorm.io/gorm" ) // Global variables var DB *gorm.DB type Data struct { T *int `json:"TYPE"` // TYPE Request *int `json:"REQUEST"` // REQUEST TYPE Result *int `json:"RESPONSE,omitempty"` // RESPONSE TYPE Hostname *string `json:"HOSTNAME"` // HOSTNAME Data *map[string]any `json:"DATA,omitempty"` // DATA } func handleInitialConnection(conn net.Conn) { initial_request := make([]byte, 256) n, err := conn.Read(initial_request) if err == io.EOF { fmt.Fprintf(os.Stderr, "ERROR: client has disconnected: %s\n", conn.RemoteAddr()) conn.Close() return } else if err != nil { fmt.Fprintf(os.Stderr, "ERROR: %s failed with the error %v, closing connection\n", conn.RemoteAddr(), err) conn.Close() return } if n == 0 { fmt.Fprintf(os.Stderr, "ERROR: received zero bytes, closing connection\n") conn.Close() return } data := Data{} err = json.Unmarshal(initial_request[:n], &data) if err != nil { fmt.Fprintf(os.Stderr, "ERROR: failed to parse data in initial request, closing connection\n") conn.Close() return } if *data.T != 0 || *data.Request != 0 || len(*data.Hostname) < 1 { fmt.Fprintf(os.Stderr, "data: %v\n", data) fmt.Fprintf(os.Stderr, "data.Hostname: \"%s\"\n", data.Hostname) fmt.Fprintf(os.Stderr, "ERROR: invalid initial request, closing connection from %s\n", conn.RemoteAddr()) conn.Close() return } switch *data.Request { case 0: var connectedResponse Data connectedResponse.T = new(int) connectedResponse.Result = new(int) *connectedResponse.T = 1 foundHost := false for _, host := range GetAllHosts() { if strings.ToLower(host.Hostname) == strings.ToLower(*data.Hostname) { *connectedResponse.T = 1 *connectedResponse.Result = 0 fmt.Printf("Found host in database\n") foundHost = true break } } if !foundHost { if ok := CreateHost(*data.Hostname); !ok { fmt.Fprintf(os.Stderr, "ERROR: failed to create host in database, disconnecting client\n") conn.Close() return } fmt.Printf("Created host in database") } bytes, err := json.Marshal(connectedResponse) if err != nil { fmt.Fprintf(os.Stderr, "ERROR: failed to create JSON to send back to client, closing connection\n") conn.Close() return } conn.Write(bytes) // TODO: Handle client connections differently handleServerConnection(conn) } } func handleServerConnection(conn net.Conn) { data := make([]byte, 4028) for { n, err := conn.Read(data) if err == io.EOF { fmt.Fprintf(os.Stderr, "ERROR: client has disconnected: %s\n", conn.RemoteAddr()) conn.Close() return } else if err != nil { fmt.Fprintf(os.Stderr, "ERROR: %s failed with the error %v, closing connection\n", conn.RemoteAddr(), err) conn.Close() return } if n == 0 { continue } fmt.Printf("Read %d bytes from %s\n", n, conn.RemoteAddr()) fmt.Printf("data:\n\t'%s'\n", data) } } func main() { cfg, err := ini.Load("config.ini") if err != nil { log.Fatal("Fail to read file: ", err) } // Classic read of values, default section can be represented as empty string log.Println("App Mode: ", cfg.Section("").Key("app_mode").String()) log.Println("Database Type: ", cfg.Section("database").Key("type").String()) if strings.ToLower(cfg.Section("database").Key("type").String()) == "postgres" { log.Println("Database Host: ", cfg.Section("database").Key("host").String()) log.Println("Database Port: ", cfg.Section("database").Key("port").String()) log.Println("Database Username:", cfg.Section("database").Key("username").String()) log.Println("Database Name: ", cfg.Section("database").Key("DB_name").String()) databaseConnectString := fmt.Sprintf( "host=%s port=%s user=%s DBname=%s password=%s sslmode=disable", cfg.Section("database").Key("host"), cfg.Section("database").Key("port"), cfg.Section("database").Key("user"), cfg.Section("database").Key("DB_name"), cfg.Section("database").Key("password"), ) log.Println(databaseConnectString) DB, err = gorm.Open(postgres.Open(databaseConnectString), &gorm.Config{TranslateError: true}) if err != nil { log.Fatal("Failed to connect to PostgreSQL DB") } } else if strings.ToLower(cfg.Section("database").Key("type").String()) == "sqlite" { sqlite_path := cfg.Section("database").Key("path").String() log.Println("sqlite DB Path: ", sqlite_path) DB, err = gorm.Open(sqlite.Open(sqlite_path), &gorm.Config{TranslateError: true}) if err != nil { log.Fatal("ERROR: failed to open sqlite database") } } port := cfg.Section("connection").Key("port").String() log.Println("Successfully connected to database") err = DB.AutoMigrate(&CPU{}) if err != nil { log.Fatal("Failed to migrate CPU schema") } err = DB.AutoMigrate(&Memory{}) if err != nil { log.Fatal("Failed to migrate Memory schema") } err = DB.AutoMigrate(&Host{}) if err != nil { log.Fatal("Failed to migrate Host schema") } addr := fmt.Sprintf(":%s", port) ln, err := net.Listen("tcp", addr) if err != nil { fmt.Fprintln(os.Stderr, "ERROR: failed to listen on port with error: ", err) os.Exit(-1) } fmt.Printf("Listening on %s\n", addr) for { conn, err := ln.Accept() if err != nil { fmt.Fprintln(os.Stderr, "ERROR: failed to connect the error: ", err) continue } go handleInitialConnection(conn) } }