server/main.go

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)
	}
}