feat: Add flag verification function

This commit is contained in:
2025-04-04 15:35:46 +02:00
parent 6651725df5
commit 4e0ecf0ad1

View File

@@ -1,13 +1,46 @@
from pwn import log from pwn import log
import re import re
import sys 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): def extract_flag(data, regex):
match = re.findall(regex, data) match = re.findall(regex, data)
if match: if match:
log.success(f"Flag found: {match[0]}") log.success(f"Flag found: {match[0]}")
sys.exit(0) return match[0]
else: else:
log.failure("Flag not found.") log.failure("Flag not found.")
sys.exit(1) sys.exit(1)