From 75ae5feb0956713e4c6cac375a758e6496cf9296 Mon Sep 17 00:00:00 2001 From: mayankmittal-iitr Date: Sun, 19 Dec 2021 00:09:36 +0530 Subject: [PATCH] flag Sumission service --- lib/mongo/constants.go | 10 +++++ lib/mongo/read.go | 20 +++++++++ services/flaghandlerservice/controller.go | 43 ++++++++++++++++++ services/flaghandlerservice/errors.go | 33 ++++++++++++++ services/flaghandlerservice/server.go | 54 ++++++++++------------- types/mongo.go | 2 +- 6 files changed, 130 insertions(+), 32 deletions(-) create mode 100644 services/flaghandlerservice/errors.go diff --git a/lib/mongo/constants.go b/lib/mongo/constants.go index 33104a75..6a87fb8c 100644 --- a/lib/mongo/constants.go +++ b/lib/mongo/constants.go @@ -10,4 +10,14 @@ const ( UsernameKey = "username" PasswordKey = "password" + + FlagsCollection = "flags" + + ChallengesCollection = "challenges" + + SubmissionsCollection = "submissions" + + ValueKey = "value" + + IDKey = "id" ) diff --git a/lib/mongo/read.go b/lib/mongo/read.go index d70717b8..facfe1e8 100644 --- a/lib/mongo/read.go +++ b/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 +} diff --git a/services/flaghandlerservice/controller.go b/services/flaghandlerservice/controller.go index 11aad807..a938c08a 100644 --- a/services/flaghandlerservice/controller.go +++ b/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 +} diff --git a/services/flaghandlerservice/errors.go b/services/flaghandlerservice/errors.go new file mode 100644 index 00000000..1da7a7f2 --- /dev/null +++ b/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 `\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" +) diff --git a/services/flaghandlerservice/server.go b/services/flaghandlerservice/server.go index 0b79d502..60f5301f 100644 --- a/services/flaghandlerservice/server.go +++ b/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 `\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 } } diff --git a/types/mongo.go b/types/mongo.go index 49a8c493..097904ed 100644 --- a/types/mongo.go +++ b/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"`