flag Sumission service

lib/mongo/constants.go

@@ -10,4 +10,14 @@ const (
 	UsernameKey = "username"
 
 	PasswordKey = "password"
+
+	FlagsCollection = "flags"
+
+	ChallengesCollection = "challenges"
+
+	SubmissionsCollection = "submissions"
+
+	ValueKey = "value"
+
+	IDKey = "id"
 )

lib/mongo/read.go

@@ -43,3 +43,23 @@ func FetchSingleTeam(teamName string) (*types.CTFTeam, error) {
 	}
 	return team, nil
 }
+
+func FetchFlag(flagValue string) (*types.Flag, error) {
+	collection := link.Collection(FlagsCollection)
+	flag := &types.Flag{}
+	ctx := context.Background()
+	if err := collection.FindOne(ctx, bson.M{ValueKey: flagValue}).Decode(flag); err != nil {
+		return nil, err
+	}
+	return flag, nil
+}
+
+func FetchChallenge(challengeId int) (*types.Challenge, error) {
+	collection := link.Collection(ChallengesCollection)
+	challenge := &types.Challenge{}
+	ctx := context.Background()
+	if err := collection.FindOne(ctx, bson.M{IDKey: challengeId}).Decode(challenge); err != nil {
+		return nil, err
+	}
+	return challenge, nil
+}

services/flaghandlerservice/controller.go

@@ -1 +1,44 @@
 package flaghandlerservice
+
+import (
+	"context"
+	"log"
+
+	"github.com/sdslabs/katana/lib/mongo"
+	"github.com/sdslabs/katana/types"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo/options"
+)
+
+func submitFlag(value string, team types.CTFTeam) (bool, int) {
+	flag, err := mongo.FetchFlag(value)
+	var points int
+	if err != nil {
+		if flag.TeamID == team.Index {
+			return false, 0
+		}
+		submission := &types.Submission{}
+		submission.ChallengeID = flag.ChallengeID
+		submission.Flag = flag.Value
+		submission.Submitter = team.Index //time of submission could also be stored
+		if res, err := mongo.InsertOne(context.Background(), mongo.SubmissionsCollection, submission); err != nil {
+			log.Println(err)
+			return false, 0
+		} else {
+			log.Println(res)
+		}
+		if res, err := mongo.FetchChallenge(flag.ChallengeID); err != nil {
+			log.Println(err)
+			return false, 0
+		} else {
+			team.Score += res.Points
+			points = res.Points
+		}
+		if err := mongo.UpdateOne(context.Background(), mongo.TeamsCollection, bson.M{"id": team.Index}, team, options.FindOneAndUpdate().SetUpsert(false)); err != nil {
+			log.Println(err)
+			return false, 0
+		}
+		return true, points
+	}
+	return false, 0
+}

services/flaghandlerservice/errors.go

@@ -0,0 +1,33 @@
+package flaghandlerservice
+
+const (
+	ServiceFail = "Failed to start flag submision service"
+
+	ServiceSuccess = "Flag Submission Service Started at port "
+
+	Connected = "Connected to"
+
+	ClosingError = "Failed to close\n"
+
+	ReadError = "Connection Read error\n"
+
+	InitInstruction = "Connected to Flag Submission Service\nInitiate your session by `init <teamName> <password>`\n"
+
+	InvalidCommand = "Invalid Command\n"
+
+	InvalidParams = "Invalid Login Parameters \n"
+
+	TeamAlreadyExists = "Team is already Logged in\n"
+
+	TeamConnected = "Team successfully connected,\n Enter flags to submit them\n"
+
+	InvalidCreds = "Invalid Team Credentials\n"
+
+	SubmitSuccess = "Flag submitted successfully, points: "
+
+	InvalidFlag = "Invalid FLag\n"
+
+	NoLogin = "Please login Team first \n"
+
+	WriteError = "Failed To Write"
+)

services/flaghandlerservice/server.go

@@ -13,84 +13,80 @@ import (
 	"github.com/sdslabs/katana/types"
 )
 
-type Team struct {
-	TeamName string
-	TeamID   int
-}
-
 func server() {
 	ln, err := net.Listen(configs.FlagConfig.SubmissionServicePort, "tcp")
 	if err != nil {
-		log.Fatal("Failed to Start Flag Submission Service")
+		log.Fatal(ServiceFail)
 	}
 	defer ln.Close()
-	log.Println("Flag Submission Service Started at port", configs.FlagConfig.SubmissionServicePort)
-	connectedTeam := Team{}
+	log.Println(ServiceSuccess, configs.FlagConfig.SubmissionServicePort)
+	connectedTeam := types.CTFTeam{}
 
 	for {
 		conn, err := ln.Accept()
 		if err != nil {
 			fmt.Println(err)
 			if err := conn.Close(); err != nil {
-				log.Println("Failed to close", err)
+				log.Println(ClosingError, err)
 			}
 			continue
 		}
 
-		log.Println("Connected to", conn.RemoteAddr())
+		log.Println(Connected, conn.RemoteAddr())
 		go handleConnection(conn, connectedTeam)
 	}
 }
 
-func handleConnection(conn net.Conn, connectedTeam Team) {
+func handleConnection(conn net.Conn, connectedTeam types.CTFTeam) {
 	defer func() {
 		if err := conn.Close(); err != nil {
-			log.Println("Error Closing", err)
+			log.Println(ClosingError, err)
 		}
 	}()
-	writeToCient(conn, "Connected to Flag Submission Service\nInitiate your session by `init <teamName> <password>`\n")
+	writeToCient(conn, InitInstruction)
 
 	for {
 		cmdLine := make([]byte, (1024 * 4))
 		n, err := conn.Read(cmdLine)
 
 		if n == 0 || err != nil {
-			log.Println("Connection Read err", err)
+			log.Println(ReadError, err)
 			return
 		}
 
 		cmd, param, password := parseCommand(string(cmdLine[0:n]))
 
 		if cmd == "" {
-			writeToCient(conn, "Inavlid Command\n")
+			writeToCient(conn, InvalidCommand)
 			continue
 		}
 		switch cmd {
 		case "init":
 			if param == "" || password == "" {
-				writeToCient(conn, "Invalid Login Parameters\n")
+				writeToCient(conn, InvalidParams)
 				continue
-			} else if (Team{}) != connectedTeam {
-				writeToCient(conn, "Team is already Logged in\n")
+			} else if (types.CTFTeam{}) != connectedTeam {
+				writeToCient(conn, TeamAlreadyExists)
 				continue
 			} else {
-				if checkTeam(param) {
-					connectedTeam.TeamAddress = conn.RemoteAddr().String()
-					connectedTeam.TeamID = param
-					writeToCient(conn, "Team successfully connected,\n Enter flags to submit them\n")
+				if condition, team := checkTeam(param, password); condition {
+					connectedTeam = team
+					writeToCient(conn, TeamConnected)
 					continue
 				} else {
-					writeToCient(conn, "Invalid TeamID\n")
+					writeToCient(conn, InvalidCreds)
 					continue
 				}
 			}
 		case "exit":
 
 		default:
-			if status, points := submitFlag(cmd); status {
-				writeToCient(conn, "Submitted successfully, points:"+strconv.Itoa(points)+"\n")
+			if (types.CTFTeam{}) == connectedTeam {
+				writeToCient(conn, NoLogin)
+			} else if status, points := submitFlag(cmd, connectedTeam); status {
+				writeToCient(conn, SubmitSuccess+strconv.Itoa(points)+"\n")
 			} else {
-				writeToCient(conn, "Invalid Flag")
+				writeToCient(conn, InvalidFlag)
 			}
 		}
 	}
@@ -129,13 +125,9 @@ func checkTeam(teamName string, password string) (bool, types.CTFTeam) {
 	return false, *team
 }
 
-func submitFlag(flag string) (bool, int) {
-	return true, 10
-}
-
 func writeToCient(conn net.Conn, message string) {
 	if _, err := conn.Write([]byte(message)); err != nil {
-		log.Println("failed to write", err)
+		log.Println(WriteError, err)
 		return
 	}
 }

types/mongo.go

@@ -8,7 +8,7 @@ type AdminUser struct {
 }
 
 type CTFTeam struct {
-	Index     int    `json:"id" bson:"password" binding:"required"`
+	Index     int    `json:"id" bson:"id" binding:"required"`
 	Name      string `json:"name" bson:"username" binding:"required"`
 	PodName   string `json:"podname" bson:"podname" binding:"required"`
 	Password  string `json:"password" bson:"password" binding:"required"`