223 lines
7.5 KiB
Python
223 lines
7.5 KiB
Python
# src/ctf/utils.py
|
|
# basic utilities
|
|
|
|
import toml
|
|
from pathlib import Path
|
|
from platformdirs import user_config_dir
|
|
# Parse the config file, returning a config dictionary with relevant config options
|
|
|
|
|
|
# Load the config from file and parse with TOML
|
|
def load_config(config = f"{user_config_dir()}/ctf-config.toml") -> dict:
|
|
p = Path(config)
|
|
if p.exists():
|
|
return toml.load(p)
|
|
return{}
|
|
|
|
# Write a dictionary to the config file
|
|
def write_config(data: dict, config = f"{user_config_dir()}/ctf"):
|
|
with open(config, "w") as f:
|
|
toml.dump(data, f)
|
|
|
|
# return a list of path objects for each catagory in a competition
|
|
def active_categories(p: Path) -> list:
|
|
cats = []
|
|
if not p.exists(): return []
|
|
for cat in p.iterdir():
|
|
cats.append(cat)
|
|
return cats
|
|
|
|
#return a lsit of Path objects of competitions in the base directory
|
|
def active_competitions(dir: str) -> dict:
|
|
comps = {}
|
|
p = Path(dir)
|
|
if not p.exists(): return {}
|
|
for item in p.iterdir():
|
|
if item.name == "tools": continue
|
|
comps[item] = active_categories(item)
|
|
print(item.name)
|
|
print(comps[item])
|
|
return comps
|
|
|
|
# {{{ suggest_patterns
|
|
def suggest_patterns(s: str, comp_name: str = "") -> list[str]:
|
|
"""Suggests a list of regex patterns from a sample flag string, optionally incorporating the competition name."""
|
|
if not s or not s.strip():
|
|
raise ValueError("Input string cannot be empty or whitespace only.")
|
|
|
|
def get_type(c):
|
|
if c.isupper():
|
|
return 'U'
|
|
elif c.islower():
|
|
return 'L'
|
|
elif c.isdigit():
|
|
return 'D'
|
|
else:
|
|
return 'S'
|
|
|
|
groups = []
|
|
comp_idx = -1
|
|
matched_prefix = ""
|
|
if comp_name and comp_name.strip():
|
|
comp_idx = s.lower().find(comp_name.lower())
|
|
|
|
if comp_idx != -1:
|
|
matched_prefix = s[comp_idx : comp_idx + len(comp_name)]
|
|
# Parse prefix before competition name
|
|
prefix_part = s[:comp_idx]
|
|
if prefix_part:
|
|
current_type = None
|
|
current_chars = []
|
|
for c in prefix_part:
|
|
ctype = get_type(c)
|
|
if ctype != current_type:
|
|
if current_type is not None:
|
|
groups.append((current_type, "".join(current_chars)))
|
|
current_type = ctype
|
|
current_chars = [c]
|
|
else:
|
|
current_chars.append(c)
|
|
if current_type is not None:
|
|
groups.append((current_type, "".join(current_chars)))
|
|
|
|
# Insert competition name group
|
|
groups.append(('C', matched_prefix))
|
|
|
|
# Parse suffix after competition name
|
|
suffix_part = s[comp_idx + len(comp_name):]
|
|
if suffix_part:
|
|
current_type = None
|
|
current_chars = []
|
|
for c in suffix_part:
|
|
ctype = get_type(c)
|
|
if ctype != current_type:
|
|
if current_type is not None:
|
|
groups.append((current_type, "".join(current_chars)))
|
|
current_type = ctype
|
|
current_chars = [c]
|
|
else:
|
|
current_chars.append(c)
|
|
if current_type is not None:
|
|
groups.append((current_type, "".join(current_chars)))
|
|
else:
|
|
# Normal grouping without competition name
|
|
current_type = None
|
|
current_chars = []
|
|
for c in s:
|
|
ctype = get_type(c)
|
|
if ctype != current_type:
|
|
if current_type is not None:
|
|
groups.append((current_type, "".join(current_chars)))
|
|
current_type = ctype
|
|
current_chars = [c]
|
|
else:
|
|
current_chars.append(c)
|
|
if current_type is not None:
|
|
groups.append((current_type, "".join(current_chars)))
|
|
|
|
def escape_special(chars):
|
|
res = []
|
|
for c in chars:
|
|
if c in r".+*?^$()[]{}|\/":
|
|
res.append("\\" + c)
|
|
else:
|
|
res.append(c)
|
|
return "".join(res)
|
|
|
|
def build_pat(use_counts=True, merge_case=False):
|
|
parts = []
|
|
for gtype, gchars in groups:
|
|
if gtype == 'C':
|
|
parts.append(escape_special(gchars))
|
|
elif gtype == 'U':
|
|
if merge_case:
|
|
parts.append(f"[A-Za-z]{{{len(gchars)}}}" if use_counts else "[A-Za-z]+")
|
|
else:
|
|
parts.append(f"[A-Z]{{{len(gchars)}}}" if use_counts else "[A-Z]+")
|
|
elif gtype == 'L':
|
|
if merge_case:
|
|
parts.append(f"[A-Za-z]{{{len(gchars)}}}" if use_counts else "[A-Za-z]+")
|
|
else:
|
|
parts.append(f"[a-z]{{{len(gchars)}}}" if use_counts else "[a-z]+")
|
|
elif gtype == 'D':
|
|
parts.append(f"\\d{{{len(gchars)}}}" if use_counts else "\\d+")
|
|
else:
|
|
parts.append(escape_special(gchars))
|
|
return "".join(parts)
|
|
|
|
candidates = []
|
|
|
|
# 1. Unanchored patterns (ideal for searching/scanning)
|
|
candidates.append(build_pat(use_counts=True, merge_case=False))
|
|
candidates.append(build_pat(use_counts=False, merge_case=False))
|
|
candidates.append(build_pat(use_counts=True, merge_case=True))
|
|
candidates.append(build_pat(use_counts=False, merge_case=True))
|
|
|
|
# Custom character class
|
|
unique_specials = set()
|
|
has_alpha = False
|
|
has_digit = False
|
|
for c in s:
|
|
if c.isalnum():
|
|
if c.isalpha():
|
|
has_alpha = True
|
|
if c.isdigit():
|
|
has_digit = True
|
|
else:
|
|
unique_specials.add(c)
|
|
char_class_parts = []
|
|
if has_alpha:
|
|
char_class_parts.append("A-Za-z")
|
|
if has_digit:
|
|
char_class_parts.append("0-9")
|
|
for spec in sorted(unique_specials):
|
|
if spec in r"-[]\^$":
|
|
char_class_parts.append("\\" + spec)
|
|
else:
|
|
char_class_parts.append(spec)
|
|
candidates.append(f"[{''.join(char_class_parts)}]+")
|
|
candidates.append(escape_special(s))
|
|
|
|
# Wildcard patterns with .* and .*? (unanchored)
|
|
if comp_idx == 0:
|
|
suffix_part = s[len(comp_name):]
|
|
esc_prefix = escape_special(matched_prefix)
|
|
if suffix_part.startswith("{") and suffix_part.endswith("}"):
|
|
candidates.append(f"{esc_prefix}\\{{.*\\}}")
|
|
candidates.append(f"{esc_prefix}\\{{.*?\\}}")
|
|
candidates.append(f"{esc_prefix}\\{{[^}}]*\\}}")
|
|
candidates.append(f"{esc_prefix}\\{{[A-Za-z0-9_\\-]+\\}}")
|
|
inner_content = suffix_part[1:-1]
|
|
if all(c.islower() or c.isdigit() or c == '_' for c in inner_content):
|
|
candidates.append(f"{esc_prefix}\\{{[a-z0-9_]+\\}}")
|
|
else:
|
|
candidates.append(f"{esc_prefix}.*")
|
|
candidates.append(f"{esc_prefix}.*?")
|
|
else:
|
|
candidates.append(".*")
|
|
candidates.append(".*?")
|
|
|
|
# 2. Anchored versions of the above (ideal for strict validation)
|
|
anchored_candidates = []
|
|
for cand in candidates:
|
|
anchored_candidates.append(f"^{cand}$")
|
|
candidates.extend(anchored_candidates)
|
|
|
|
# Deduplicate while preserving order
|
|
seen = set()
|
|
result = []
|
|
for c in candidates:
|
|
if c not in seen:
|
|
seen.add(c)
|
|
result.append(c)
|
|
return result
|
|
# }}}
|
|
|
|
# Load variables to export
|
|
config = load_config("/home/venus/code/ctf/config.toml")
|
|
competition = config["Competition"]
|
|
enviroment = config["Enviroment"]
|
|
# base_dir = config["ctf_dir"]
|
|
|
|
|