admin-dashboard-backend/socketserver/hub.go

package socketserver

import (
	"encoding/json"
	"fmt"
	"janex/admin-dashboard-backend/modules/cache"
	"janex/admin-dashboard-backend/modules/database"
	"janex/admin-dashboard-backend/modules/grouptasks"
	"janex/admin-dashboard-backend/modules/structs"
	"janex/admin-dashboard-backend/modules/utils"
	"janex/admin-dashboard-backend/socketclients"
	"time"

	"github.com/gofiber/websocket/v2"
	"github.com/google/uuid"
	"github.com/rs/zerolog/log"
)

var register = make(chan *structs.SocketClient)
var broadcast = make(chan structs.SocketMessage)
var unregister = make(chan *websocket.Conn)

func RunHub() {
	for {
		select {
		case newSocketClient := <-register:
			userId := fmt.Sprintf("%v", newSocketClient.Conn.Locals("userId"))
			sessionId := fmt.Sprintf("%v", newSocketClient.Conn.Locals("sessionId"))

			// close connection instantly if sessionId is empty
			if sessionId == "<nil>" {
				newSocketClient.SendUnauthorizedCloseMessage()
				continue
			}

			newSocketClient.SessionId = sessionId
			newSocketClient.UserId = userId

			cache.AddSocketClient(sessionId, newSocketClient)

			// check that user session is not expired
			var userSession structs.UserSession

			database.DB.First(&userSession, "id = ?", sessionId)

			if userSession.Id != "" && time.Now().After(userSession.ExpiresAt) {
				newSocketClient.SendUnauthorizedCloseMessage()
				database.DB.Delete(&structs.UserSession{}, "id = ?", sessionId)
				continue
			}

			// update session last used time
			database.DB.Model(&structs.UserSession{}).Where("id = ?", sessionId).Updates(structs.UserSession{
				LastUsed:  time.Now(),
				ExpiresAt: utils.GetSessionExpiresAtTime(),
			})

			// init data message
			var user structs.User

			database.DB.First(&user, "id = ?", userId)

			newSocketClient.SendMessage(structs.SendSocketMessage{
				Cmd: utils.SentCmdInitUserSocketConnection,
				Body: structs.InitUserSocketConnection{
					User: structs.UserData{
						Id:          user.Id,
						RoleId:      user.RoleId,
						Username:    user.Username,
						Email:       user.Email,
						Sessions:    socketclients.GetUserSessions(userId),
						Permissions: socketclients.GetPermissionsByRoleId(user.RoleId),
					},
					CategoryGroups:  cache.GetCategoryGroupsSorted(),
					GroupTasks:      grouptasks.GetAllGroupTasks(),
					GroupTasksSteps: grouptasks.GetAllGroupTasksSteps(),
					AllUsers:        socketclients.GetAllUsers(),
					Scanners:        socketclients.GetAllScanners(),
					AllRoles:        socketclients.GetAllRoles(),
					AdminArea: structs.InitUserSocketConnectionAdminArea{
						RolesPermissions: socketclients.GetAdminAreaRolesPermissions(),
					},
				},
			})

			socketclients.UpdateConnectedUsers(userId)
			socketclients.UpdateUserSessionsForUser(userId, sessionId)

		case data := <-broadcast:
			var receivedMessage structs.ReceivedMessage

			if err := json.Unmarshal(data.Msg, &receivedMessage); err != nil {
				log.Error().Msgf("Failed to unmarshal received msg, err: %s", err)
				continue
			}

			log.Info().Msgf("Received message: %v %v", receivedMessage, receivedMessage.Cmd)

			switch receivedMessage.Cmd {
			case utils.ReceivedCmdStartGroupTasks:
				category := receivedMessage.Body["category"].(string)
				groupId := receivedMessage.Body["id"].(string)

				globalInputsJsonString := utils.MarshalJson(receivedMessage.Body["globalInputs"])

				groupTaskId := uuid.New().String()

				groupTasks := &structs.GroupTasks{
					Id:               groupTaskId,
					CreatorUserId:    data.Conn.Locals("userId").(string),
					Category:         category,
					GroupId:          groupId,
					GroupName:        receivedMessage.Body["groupName"].(string),
					Description:      receivedMessage.Body["description"].(string),
					CurrentTasksStep: 1,
					NumberOfSteps:    uint8(receivedMessage.Body["numberOfSteps"].(float64)),
					Status:           utils.GroupTasksStatusRunning,
					GlobalInputs:     globalInputsJsonString,
					StartedAt:        time.Now(),
					RememberId:       receivedMessage.Body["rememberId"].(string),
				}

				database.DB.Create(groupTasks)

				socketclients.BroadcastMessage(structs.SendSocketMessage{
					Cmd:  utils.SentCmdNewGroupTaskStarted,
					Body: groupTasks,
				})

				go grouptasks.RunGroupTask(grouptasks.RunGroupTaskArgs{
					CreatorUserId: data.Conn.Locals("userId").(string),
					StartType:     grouptasks.RunGroupTaskStartTypeNormal,
					GroupTaskId:   groupTaskId,
					Category:      category,
					GroupId:       groupId,
					Step:          1,
					TaskStepId:    "",
					GlobalInputs:  globalInputsJsonString,
				})
				break
			case utils.ReceivedCmdTaskFailedTryAgainRunTaskStep:
				go grouptasks.RunGroupTask(grouptasks.RunGroupTaskArgs{
					CreatorUserId: data.Conn.Locals("userId").(string),
					StartType:     grouptasks.RunGroupTaskStartTypeTryAgain,
					GroupTaskId:   receivedMessage.Body["groupTaskId"].(string),
					Category:      receivedMessage.Body["category"].(string),
					GroupId:       receivedMessage.Body["groupId"].(string),
					Step:          uint8(receivedMessage.Body["step"].(float64)),
					TaskStepId:    receivedMessage.Body["taskStepId"].(string),
				})
				break
			case utils.ReceivedCmdTaskContinueTaskStep:
				go grouptasks.RunGroupTask(grouptasks.RunGroupTaskArgs{
					CreatorUserId: data.Conn.Locals("userId").(string),
					StartType:     grouptasks.RunGroupTaskStartTypeTryAgain,
					GroupTaskId:   receivedMessage.Body["groupTaskId"].(string),
					Category:      receivedMessage.Body["category"].(string),
					GroupId:       receivedMessage.Body["groupId"].(string),
					Step:          uint8(receivedMessage.Body["step"].(float64)),
					TaskStepId:    receivedMessage.Body["taskStepId"].(string),
					TaskInputs:    utils.MarshalJson(receivedMessage.Body["taskInputs"]),
				})
				break
			case utils.ReceivedCmdReloadGroupTasks:
				category := receivedMessage.Body["category"].(string)

				socketclients.BroadcastMessage(structs.SendSocketMessage{
					Cmd:  utils.SentCmdReloadingGroupTasks,
					Body: category,
				})

				grouptasks.LoadGroups(category)
				break
			case utils.ReceivedCmdTaskLocking:
				cache.AddLockedGroupTaskStep(structs.LockedGroupTaskSteps{
					LockedByUserId: receivedMessage.Body["lockedByUserId"].(string),
					GroupTaskId:    receivedMessage.Body["groupTaskId"].(string),
					Step:           uint8(receivedMessage.Body["step"].(float64)),
					LockedAt:       time.Now(),
					RememberId:     receivedMessage.Body["rememberId"].(string),
				})

				socketclients.BroadcastMessage(structs.SendSocketMessage{
					Cmd:  utils.SentCmdTaskLocked,
					Body: receivedMessage.Body,
				})

				cache.AddGroupTaskStepsInput(structs.GroupTaskStepsInput{
					GroupTaskId:   receivedMessage.Body["groupTaskId"].(string),
					Step:          uint8(receivedMessage.Body["step"].(float64)),
					ParameterName: receivedMessage.Body["parameterName"].(string),
					Value:         receivedMessage.Body["value"],
				})
				break
			case utils.ReceivedCmdUpdateUserProfile:
				socketclients.UpdateUserProfile(data.Conn, receivedMessage.Body["changes"].(map[string]interface{}))
				break
			case utils.ReceivedCmdAdminAreaCreateNewRole:
				log.Debug().Msgf("create new role")

				// TODO: check permissions

				role := structs.Role{
					Id:           uuid.New().String(),
					DisplayName:  "New Role " + uuid.New().String(),
					Description:  "Role description",
					CreatedAt:    time.Now(),
					SortingOrder: database.GetRoleSortingOrder(),
				}

				database.DB.Create(&role)

				socketclients.BroadcastMessage(structs.SendSocketMessage{
					Cmd:  utils.SentCmdAdminAreaNewRoleCreated,
					Body: role,
				})
				break
			case utils.ReceivedCmdAdminAreaUpdateRole:
				socketclients.AdminAreaUpdateRole(data.Conn.Locals("sessionId").(string), receivedMessage.Body)
				break
			default:
				log.Error().Msgf("Received unknown message: %v", receivedMessage)
				break
			}

		case connection := <-unregister:
			cache.DeleteClientByConn(connection)

			if connection.Locals("userId") != nil && connection.Locals("sessionId") != nil {
				userId := connection.Locals("userId").(string)
				sessionId := connection.Locals("sessionId").(string)

				database.DB.Model(&structs.User{}).Where("id = ?", userId).Updates(structs.User{
					LastOnline: time.Now(),
				})

				socketclients.UpdateUserSessionsForUser(userId, sessionId)

				socketclients.UpdateConnectedUsers(userId)
			}
		}
	}
}