func TestFlagExist(t *testing.T) { db, err := openDB() defer db.Close() flg := steward.Flag{ID: 0, Flag: "tralala", Round: 5, TeamID: 10, ServiceID: 4, Cred: "1:2"} err = steward.AddFlag(db.db, flg) exist, err := steward.FlagExist(db.db, flg.Flag) if !exist { log.Fatalln("Exist flag does not exist:", err) } exist, err = steward.FlagExist(db.db, "not_exist_flag") if exist { log.Fatalln("Not exist flag is exist:", err) } }
func handler(conn net.Conn, db *sql.DB, priv *rsa.PrivateKey, attackFlow chan scoreboard.Attack) { addr := conn.RemoteAddr().String() defer conn.Close() fmt.Fprint(conn, greetingMsg) flag, err := bufio.NewReader(conn).ReadString('\n') if err != nil { log.Println("Read error:", err) } flag = strings.Trim(flag, "\n") log.Printf("\tGet flag %s from %s", flag, addr) valid, err := vexillary.ValidFlag(flag, priv.PublicKey) if err != nil { log.Println("\tValidate flag failed:", err) } if !valid { fmt.Fprint(conn, invalidFlagMsg) return } exist, err := steward.FlagExist(db, flag) if err != nil { log.Println("\tExist flag check failed:", err) fmt.Fprint(conn, internalErrorMsg) return } if !exist { fmt.Fprint(conn, flagDoesNotExistMsg) return } flg, err := steward.GetFlagInfo(db, flag) if err != nil { log.Println("\tGet flag info failed:", err) fmt.Fprint(conn, internalErrorMsg) return } captured, err := steward.AlreadyCaptured(db, flg.ID) if err != nil { log.Println("\tAlready captured check failed:", err) fmt.Fprint(conn, internalErrorMsg) return } if captured { fmt.Fprint(conn, alreadyCapturedMsg) return } team, err := teamByAddr(db, addr) if err != nil { log.Println("\tGet team by ip failed:", err) fmt.Fprint(conn, invalidTeamMsg) return } if flg.TeamID == team.ID { log.Printf("\tTeam %s try to send their flag", team.Name) fmt.Fprint(conn, flagYoursMsg) return } round, err := steward.CurrentRound(db) if round.ID != flg.Round { log.Printf("\t%s try to send flag from past round", team.Name) fmt.Fprint(conn, flagExpiredMsg) return } roundEndTime := round.StartTime.Add(round.Len) if time.Now().After(roundEndTime) { log.Printf("\t%s try to send flag from finished round", team.Name) fmt.Fprint(conn, flagExpiredMsg) return } halfStatus := steward.Status{flg.Round, team.ID, flg.ServiceID, steward.StatusUnknown} state, err := steward.GetState(db, halfStatus) if state != steward.StatusUP { log.Printf("\t%s service not ok, cannot capture", team.Name) fmt.Fprint(conn, serviceNotUpMsg) return } err = steward.CaptureFlag(db, flg.ID, team.ID) if err != nil { log.Println("\tCapture flag failed:", err) fmt.Fprint(conn, internalErrorMsg) return } go func() { attack := scoreboard.Attack{ Attacker: team.ID, Victim: flg.TeamID, Service: flg.ServiceID, Timestamp: time.Now().Unix(), } select { case attackFlow <- attack: default: _ = <-attackFlow attackFlow <- attack } }() fmt.Fprint(conn, capturedMsg) }