challenge registration fixes

pwncollege · Nov 23, 2023 · 805d8fd · 805d8fd
1 parent cb72999
commit 805d8fd
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/pwnshop/__init__.py b/pwnshop/__init__.py
@@ -8,7 +8,7 @@
 
 def register_challenge(challenge):
     module = inspect.getmodule(challenge)
-    module_name = module.__name__
+    module_name = module.__name__.split(".")[-1]
     challenge_name = challenge.__name__
 
     if challenge_name in ALL_CHALLENGES:
@@ -17,3 +17,7 @@ def register_challenge(challenge):
     ALL_CHALLENGES[challenge_name] = challenge
     MODULE_LEVELS.setdefault(module_name, []).append(challenge)
     ALL_MODULES[module_name] = module
+
+def register_challenges(challenges):
+    for c in challenges:
+        register_challenge(c)