From 4e0ecf0ad1a4dad58b8cceca2180a1f423045011 Mon Sep 17 00:00:00 2001
From: Matte23 <matteo@underdesk.net>
Date: Fri, 4 Apr 2025 15:35:46 +0200
Subject: [PATCH] feat: Add flag verification function

---
 src/solvertools/flags.py | 35 ++++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/solvertools/flags.py b/src/solvertools/flags.py
index 3416f3c..21d2bde 100644
--- a/src/solvertools/flags.py
+++ b/src/solvertools/flags.py
@@ -1,13 +1,46 @@
 from pwn import log
 import re
 import sys
+import os
+import yaml
+
+
+def verify_flag(data, regex):
+    """
+    Verify the flag using challenge.yml if available. 
+    Otherwise, use regex-based verification.
+    """
+    flag = extract_flag(data, regex)
+    challenge_file = "challenge.yml"
+
+    if not os.path.exists(challenge_file):
+        # No challenge.yml, skip verification
+        log.warning("Missing challenge.yml, skipping verification.")
+        return
+
+    # challenge.yml exists, load and check against listed flags
+    try:
+        with open(challenge_file, 'r') as f:
+            data = yaml.safe_load(f)
+    except Exception as e:
+        log.failure(f"Error reading challenge.yml: {e}")
+        sys.exit(1)
+
+    flags = data.get("flags", [])
+
+    if flag in flags:
+        log.success("Flag verified successfully.")
+        sys.exit(0)
+    else:
+        log.failure("Flag not present in challenge.yml.")
+        sys.exit(2)
 
 
 def extract_flag(data, regex):
     match = re.findall(regex, data)
     if match:
         log.success(f"Flag found: {match[0]}")
-        sys.exit(0)
+        return match[0]
     else:
         log.failure("Flag not found.")
         sys.exit(1)