feat: Add flag verification function
This commit is contained in:
@@ -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)
|
||||||
|
|||||||
Reference in New Issue
Block a user