Merge pull request #29 from nsec/added-path

Reworked the CTF root directory and added the location global argument.

Author: Res260

ctf/__init__.py

@@ -6,11 +6,6 @@
 import urllib.request
 
 from ctf.logger import LOG
-from ctf.utils import (
-    find_ctf_root_directory,
-    get_ctf_script_schemas_directory,
-    get_ctf_script_templates_directory,
-)
 
 VERSION = importlib.metadata.version("ctf-script")
 
@@ -23,17 +18,6 @@
 for k, v in os.environ.items():
     ENV[k] = v
 
-match sys.argv[1] if len(sys.argv) > 1 else "":
-    case "init":
-        CTF_ROOT_DIRECTORY = os.path.join(os.getcwd(), ".")
-    case "version":
-        CTF_ROOT_DIRECTORY = ""
-    case _:
-        CTF_ROOT_DIRECTORY = find_ctf_root_directory()
-
-TEMPLATES_ROOT_DIRECTORY = get_ctf_script_templates_directory()
-SCHEMAS_ROOT_DIRECTORY = get_ctf_script_schemas_directory()
-
 
 def check_tool_version() -> None:
     with urllib.request.urlopen(

ctf/__main__.py

@@ -6,10 +6,7 @@
 from typer import Typer
 from typing_extensions import Annotated
 
-from ctf import (
-    CTF_ROOT_DIRECTORY,
-    LOG,
-)
+from ctf import ENV, LOG
 from ctf.check import app as check_app
 from ctf.deploy import app as deploy_app
 from ctf.destroy import app as destroy_app
@@ -21,6 +18,7 @@
 from ctf.redeploy import app as redeploy_app
 from ctf.services import app as services_app
 from ctf.stats import app as stats_app
+from ctf.utils import find_ctf_root_directory
 from ctf.validate import app as validate_app
 from ctf.version import app as version_app
 
@@ -44,6 +42,9 @@
 
 @app.callback()
 def global_options(
+    location: Annotated[
+        str, typer.Option("--location", help="CTF root directory location.")
+    ] = "",
     verbose: Annotated[
         bool, typer.Option("--verbose", "-v", help="Enable DEBUG logging.")
     ] = False,
@@ -52,6 +53,9 @@ def global_options(
         LOG.setLevel(logging.DEBUG)
         LOG.handlers[0].setLevel(logging.DEBUG)
 
+    if location:
+        ENV["CTF_ROOT_DIR"] = location
+
 
 def main():
     app()
@@ -61,7 +65,9 @@ def main():
     import sys
 
     if "version" not in sys.argv and "init" not in sys.argv:
-        if not os.path.isdir(s=(p := os.path.join(CTF_ROOT_DIRECTORY, "challenges"))):
+        if not os.path.isdir(
+            s=(p := os.path.join(find_ctf_root_directory(), "challenges"))
+        ):
             LOG.error(
                 msg=f"Directory `{p}` not found. Make sure this script is ran from the root directory OR set the CTF_ROOT_DIR environment variable to the root directory."
             )

ctf/check.py

@@ -4,10 +4,10 @@
 import typer
 from typing_extensions import Annotated
 
-from ctf import CTF_ROOT_DIRECTORY, ENV
+from ctf import ENV
 from ctf.generate import generate
 from ctf.logger import LOG
-from ctf.utils import check_git_lfs, terraform_binary
+from ctf.utils import check_git_lfs, find_ctf_root_directory, terraform_binary
 
 app = typer.Typer()
 
@@ -40,7 +40,7 @@ def check(
     # Then run terraform plan.
     subprocess.run(
         args=[terraform_binary(), "plan"],
-        cwd=os.path.join(CTF_ROOT_DIRECTORY, ".deploy"),
+        cwd=os.path.join(find_ctf_root_directory(), ".deploy"),
         check=True,
     )
 

ctf/deploy.py

@@ -7,13 +7,14 @@
 import typer
 from typing_extensions import Annotated
 
-from ctf import CTF_ROOT_DIRECTORY, ENV
+from ctf import ENV
 from ctf.destroy import destroy
 from ctf.generate import generate
 from ctf.logger import LOG
 from ctf.utils import (
     add_tracks_to_terraform_modules,
     check_git_lfs,
+    find_ctf_root_directory,
     get_all_available_tracks,
     get_terraform_tracks_from_modules,
     parse_track_yaml,
@@ -91,7 +92,7 @@ def deploy(
     try:
         subprocess.run(
             args=[terraform_binary(), "apply", "-auto-approve"],
-            cwd=os.path.join(CTF_ROOT_DIRECTORY, ".deploy"),
+            cwd=os.path.join(find_ctf_root_directory(), ".deploy"),
             check=True,
         )
     except subprocess.CalledProcessError:
@@ -107,7 +108,7 @@ def deploy(
 
         subprocess.run(
             args=[terraform_binary(), "apply", "-auto-approve"],
-            cwd=os.path.join(CTF_ROOT_DIRECTORY, ".deploy"),
+            cwd=os.path.join(find_ctf_root_directory(), ".deploy"),
             check=True,
         )
     except KeyboardInterrupt:
@@ -121,7 +122,9 @@ def deploy(
     for track in distinct_tracks:
         if not os.path.exists(
             path=(
-                path := os.path.join(CTF_ROOT_DIRECTORY, "challenges", track, "ansible")
+                path := os.path.join(
+                    find_ctf_root_directory(), "challenges", track, "ansible"
+                )
             )
         ):
             continue

ctf/destroy.py

@@ -5,9 +5,10 @@
 import typer
 from typing_extensions import Annotated
 
-from ctf import CTF_ROOT_DIRECTORY, ENV
+from ctf import ENV
 from ctf.logger import LOG
 from ctf.utils import (
+    find_ctf_root_directory,
     get_terraform_tracks_from_modules,
     remove_tracks_from_terraform_modules,
     terraform_binary,
@@ -50,7 +51,7 @@ def destroy(
     LOG.info(msg="tofu destroy...")
 
     if not os.path.exists(
-        path=os.path.join(CTF_ROOT_DIRECTORY, ".deploy", "modules.tf")
+        path=os.path.join(find_ctf_root_directory(), ".deploy", "modules.tf")
     ):
         LOG.critical(msg="Nothing to destroy.")
         exit(code=1)
@@ -118,7 +119,7 @@ def destroy(
                 else [f"-target=module.track-{track}" for track in terraform_tracks]
             ),
         ],
-        cwd=os.path.join(CTF_ROOT_DIRECTORY, ".deploy"),
+        cwd=os.path.join(find_ctf_root_directory(), ".deploy"),
         check=False,
     )
 

ctf/flags.py

@@ -9,9 +9,8 @@
 import yaml
 from typing_extensions import Annotated
 
-from ctf import CTF_ROOT_DIRECTORY
 from ctf.logger import LOG
-from ctf.utils import parse_track_yaml
+from ctf.utils import find_ctf_root_directory, parse_track_yaml
 
 app = typer.Typer()
 
@@ -41,7 +40,11 @@ def flags(
     distinct_tracks: set[str] = set()
 
     for entry in os.listdir(
-        path=(challenges_directory := os.path.join(CTF_ROOT_DIRECTORY, "challenges"))
+        path=(
+            challenges_directory := os.path.join(
+                find_ctf_root_directory(), "challenges"
+            )
+        )
     ):
         if os.path.isdir(
             s=(track_directory := os.path.join(challenges_directory, entry))

ctf/generate.py

@@ -4,11 +4,12 @@
 import typer
 from typing_extensions import Annotated
 
-from ctf import CTF_ROOT_DIRECTORY, ENV
+from ctf import ENV
 from ctf.logger import LOG
 from ctf.utils import (
     add_tracks_to_terraform_modules,
     create_terraform_modules_file,
+    find_ctf_root_directory,
     get_all_available_tracks,
     terraform_binary,
     validate_track_can_be_deployed,
@@ -61,10 +62,10 @@ def generate(
 
         for track in distinct_tracks:
             relpath = os.path.relpath(
-                os.path.join(CTF_ROOT_DIRECTORY, ".deploy", "common"),
+                os.path.join(find_ctf_root_directory(), ".deploy", "common"),
                 (
                     terraform_directory := os.path.join(
-                        CTF_ROOT_DIRECTORY, "challenges", track, "terraform"
+                        find_ctf_root_directory(), "challenges", track, "terraform"
                     )
                 ),
             )
@@ -103,13 +104,13 @@ def generate(
 
         subprocess.run(
             args=[terraform_binary(), "init", "-upgrade"],
-            cwd=os.path.join(CTF_ROOT_DIRECTORY, ".deploy"),
+            cwd=os.path.join(find_ctf_root_directory(), ".deploy"),
             stdout=subprocess.DEVNULL,
             check=True,
         )
         subprocess.run(
             args=[terraform_binary(), "validate"],
-            cwd=os.path.join(CTF_ROOT_DIRECTORY, ".deploy"),
+            cwd=os.path.join(find_ctf_root_directory(), ".deploy"),
             check=True,
         )
     else:

ctf/init.py

@@ -4,8 +4,9 @@
 import typer
 from typing_extensions import Annotated
 
-from ctf import CTF_ROOT_DIRECTORY, TEMPLATES_ROOT_DIRECTORY
+from ctf import ENV
 from ctf.logger import LOG
+from ctf.utils import get_ctf_script_templates_directory
 
 app = typer.Typer()
 
@@ -16,14 +17,22 @@
 def init(
     path: Annotated[
         str, typer.Argument(help="Directory in which to initialize a CTF")
-    ] = CTF_ROOT_DIRECTORY,
+    ] = "",
     force: Annotated[
         bool,
         typer.Option(
             "--force", help="Overwrite the directory if it's already initialized"
         ),
     ] = False,
 ) -> None:
+    # If path is not set, take the one from --location or CTF_ROOT_DIR, else it's the current directory.
+    if not path:
+        path = (
+            ENV.get("CTF_ROOT_DIR")
+            if "CTF_ROOT_DIR" in ENV
+            else os.path.join(os.getcwd(), ".")
+        )
+
     created_directory = False
     created_assets: list[str] = []
     try:
@@ -40,7 +49,9 @@ def init(
             )
             exit(code=1)
 
-        for asset in os.listdir(p := os.path.join(TEMPLATES_ROOT_DIRECTORY, "init")):
+        for asset in os.listdir(
+            p := os.path.join(get_ctf_script_templates_directory(), "init")
+        ):
             dst_asset = os.path.join(path, asset)
             if os.path.isdir(src_asset := os.path.join(p, asset)):
                 shutil.copytree(src_asset, dst_asset, dirs_exist_ok=True)

ctf/list.py

@@ -6,8 +6,7 @@
 from rich.table import Table
 from typing_extensions import Annotated
 
-from ctf import CTF_ROOT_DIRECTORY
-from ctf.utils import parse_post_yamls, parse_track_yaml
+from ctf.utils import find_ctf_root_directory, parse_post_yamls, parse_track_yaml
 
 app = typer.Typer()
 
@@ -23,11 +22,13 @@ def list_tracks(
     ] = ListOutputFormat.PRETTY,
 ) -> None:
     tracks: set[str] = set()
-    for track in os.listdir(path=os.path.join(CTF_ROOT_DIRECTORY, "challenges")):
+    for track in os.listdir(path=os.path.join(find_ctf_root_directory(), "challenges")):
         if os.path.isdir(
-            s=os.path.join(CTF_ROOT_DIRECTORY, "challenges", track)
+            s=os.path.join(find_ctf_root_directory(), "challenges", track)
         ) and os.path.exists(
-            path=os.path.join(CTF_ROOT_DIRECTORY, "challenges", track, "track.yaml")
+            path=os.path.join(
+                find_ctf_root_directory(), "challenges", track, "track.yaml"
+            )
         ):
             tracks.add(track)
 

ctf/new.py

@@ -8,8 +8,8 @@
 import typer
 from typing_extensions import Annotated
 
-from ctf import CTF_ROOT_DIRECTORY, TEMPLATES_ROOT_DIRECTORY
 from ctf.logger import LOG
+from ctf.utils import find_ctf_root_directory, get_ctf_script_templates_directory
 
 app = typer.Typer()
 
@@ -63,7 +63,7 @@ def new(
     if os.path.exists(
         path=(
             new_challenge_directory := os.path.join(
-                CTF_ROOT_DIRECTORY, "challenges", name
+                find_ctf_root_directory(), "challenges", name
             )
         )
     ):
@@ -82,7 +82,7 @@ def new(
 
     env = jinja2.Environment(
         loader=jinja2.FileSystemLoader(
-            searchpath=TEMPLATES_ROOT_DIRECTORY, encoding="utf-8"
+            searchpath=get_ctf_script_templates_directory(), encoding="utf-8"
         )
     )
 
@@ -191,7 +191,8 @@ def new(
     LOG.debug(msg=f"Wrote {p}.")
 
     relpath = os.path.relpath(
-        os.path.join(CTF_ROOT_DIRECTORY, ".deploy", "common"), terraform_directory
+        os.path.join(find_ctf_root_directory(), ".deploy", "common"),
+        terraform_directory,
     )
 
     os.symlink(
@@ -278,7 +279,7 @@ def new(
     if template == Template.RUST_WEBSERVICE:
         # Copy the entire challenge template
         shutil.copytree(
-            os.path.join(TEMPLATES_ROOT_DIRECTORY, "rust-webservice"),
+            os.path.join(get_ctf_script_templates_directory(), "rust-webservice"),
             ansible_challenge_directory,
             dirs_exist_ok=True,
         )