Add time limit(s)

1 files changed, 56 insertions, 0 deletions

diff --git a/src/limit.go b/src/limit.go
new file mode 100644
index 0000000..c4befac
--- /dev/null
+++ b/src/limit.go
@@ -0,0 +1,56 @@
+package main
+
+import (
+	"flag"
+	"time"
+)
+
+var (
+	endTimestamp *int64
+	endAfter *int64
+	initialLoginTime int64
+)
+
+func registerLimitFlags() {
+	endTimestamp = flag.Int64("endTimestamp", 0, "Date/Time after which flags are not accepted anymore")
+	endAfter = flag.Int64("endAfter", 0, "Seconds after the first login, after which flags are not accepted anymore")
+}
+
+func startLimitTimer() {
+	now := time.Now().Unix()
+
+	// check if endTimestamp is set
+	if *endTimestamp > 0 {
+		// Start the endTimestampTimer
+		endTimestampTimer := time.NewTimer(time.Duration(*endTimestamp - now))
+		go func() {
+			<-endTimestampTimer.C
+			// Stop all challenge containers if timer hit
+			stopAllChallengeContainers()
+		}()
+	}
+
+	// check if endAfter is set
+	if *endAfter > 0 {
+		// Start the endAfterTimer
+		endAfterTimer := time.NewTimer(time.Duration(*endAfter - now))
+		go func() {
+			<-endAfterTimer.C
+			// Stop all challenge containers if timer hit
+			stopAllChallengeContainers()
+		}()
+	}
+}
+
+// Called on every login to set the "initialLoginTime" if it's not already set
+func registerLoginForLimiter() {
+	if initialLoginTime == 0 {
+		initialLoginTime = time.Now().Unix()
+	}
+}
+
+// Called before starting containers or entering flags
+func shouldLimit() (bool) {
+	now := time.Now().Unix()
+	return (*endTimestamp > 0 && *endTimestamp < now) || (*endAfter > 0 && *endAfter < now - initialLoginTime)
+}