diff --git a/server/internal/services/servers.go b/server/internal/services/servers.go index 2ce9ff2..c1b62b5 100644 --- a/server/internal/services/servers.go +++ b/server/internal/services/servers.go @@ -79,6 +79,25 @@ func GetServerByPreRegToken(token string) (*models.Server, error) { return &s, nil } +// OSTypeFromInfo derives a coarse os_type ("windows" or "linux") from the +// agent-reported os_info string, which is formatted " ". +// Anything that is not explicitly windows defaults to linux. +func OSTypeFromInfo(osInfo string) string { + if strings.HasPrefix(strings.ToLower(osInfo), "windows") { + return "windows" + } + return "linux" +} + +// defaultConsoleFields returns the initial console configuration for a newly +// registered server based on its os_type. +func defaultConsoleFields(osType string) (protocols []string, sshPort, rdpPort int) { + if osType == "windows" { + return []string{"rdp"}, 22, 3389 + } + return []string{"ssh"}, 22, 3389 +} + func RegisterServer(serverID, preRegToken, hostname, ipAddress, osInfo string) (string, error) { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() @@ -100,18 +119,29 @@ func RegisterServer(serverID, preRegToken, hostname, ipAddress, osInfo string) ( tokenHash := HashToken(agentToken) now := time.Now() + osType := OSTypeFromInfo(osInfo) + protocols, sshPort, rdpPort := defaultConsoleFields(osType) + _, err = db.Col("servers").UpdateOne(ctx, bson.M{"server_id": serverID}, - bson.M{"$set": bson.M{ - "hostname": hostname, - "ip_address": ipAddress, - "os_info": osInfo, - "agent_token_hash": tokenHash, - "status": "active", - "last_seen": now, - "pre_reg_token": "", - "pre_reg_expires": nil, - }}, + bson.M{ + "$set": bson.M{ + "hostname": hostname, + "ip_address": ipAddress, + "os_info": osInfo, + "os_type": osType, + "agent_token_hash": tokenHash, + "status": "active", + "last_seen": now, + "pre_reg_token": "", + "pre_reg_expires": nil, + }, + "$setOnInsert": bson.M{ + "console_protocols": protocols, + "ssh_port": sshPort, + "rdp_port": rdpPort, + }, + }, ) if err != nil { return "", err diff --git a/server/internal/services/servers_console_test.go b/server/internal/services/servers_console_test.go new file mode 100644 index 0000000..1b0e973 --- /dev/null +++ b/server/internal/services/servers_console_test.go @@ -0,0 +1,18 @@ +package services + +import "testing" + +func TestOSTypeFromInfo(t *testing.T) { + cases := map[string]string{ + "windows amd64": "windows", + "linux amd64": "linux", + "linux arm64": "linux", + "": "linux", + "darwin arm64": "linux", + } + for in, want := range cases { + if got := OSTypeFromInfo(in); got != want { + t.Errorf("OSTypeFromInfo(%q) = %q, want %q", in, got, want) + } + } +}