relink.py: Check that files to process are children of CESM input data root.

Author: samsrabin

relink.py

@@ -10,6 +10,7 @@
 import argparse
 import logging
 import time
+from pathlib import Path
 
 DEFAULT_SOURCE_ROOT = "/glade/campaign/cesm/cesmdata/cseg/inputdata/"
 DEFAULT_TARGET_ROOT = (
@@ -60,18 +61,23 @@ def _handle_non_dir_entry(entry, user_uid):
     return None
 
 
-def handle_non_dir(var, user_uid):
+def handle_non_dir(var, user_uid, inputdata_root):
     """
     Check if a non-directory is owned by the user and should be processed. Passes var to a
     helper function depending on its type.
 
     Args:
         var (os.DirEntry or str): A directory entry from os.scandir(), or a string path.
         user_uid (int): The UID of the user whose files to find.
+        inputdata_root (str): The root of the directory tree containing CESM input data.
 
     Returns:
         str or None: The absolute path to the file if it's owned by the user
                      and is a regular file (not a symlink), otherwise None.
+
+    Raises:
+        TypeError: If var is not a DirEntry-like object.
+        ValueError: If the file path is not under inputdata_root.
     """
 
     # Fall back to duck typing: If var has the required DirEntry methods and members, treat it as a
@@ -80,12 +86,22 @@ def handle_non_dir(var, user_uid):
     if isinstance(var, os.DirEntry) or all(
         hasattr(var, m) for m in ["stat", "is_file", "is_symlink", "path"]
     ):
-        return _handle_non_dir_entry(var, user_uid)
+        file_path = _handle_non_dir_entry(var, user_uid)
+    else:
+        raise TypeError(
+            f"Unsure how to handle non-directory variable of type {type(var)}"
+        )
 
-    raise TypeError(f"Unsure how to handle non-directory variable of type {type(var)}")
+    # Check that resulting path is a child of inputdata_root
+    if file_path is not None and not Path(file_path).is_relative_to(inputdata_root):
+        raise ValueError(
+            f"'{file_path}' must be equivalent to or under '{inputdata_root}"
+        )
+
+    return file_path
 
 
-def find_owned_files_scandir(directory, user_uid):
+def find_owned_files_scandir(directory, user_uid, inputdata_root=DEFAULT_SOURCE_ROOT):
     """
     Efficiently find all files owned by a specific user using os.scandir().
 
@@ -95,20 +111,28 @@ def find_owned_files_scandir(directory, user_uid):
     Args:
         directory (str): The root directory to search.
         user_uid (int): The UID of the user whose files to find.
+        inputdata_root (str): The root of the directory tree containing CESM input data.
 
     Yields:
         str: Absolute paths to files owned by the user.
+
+    Raises:
+        ValueError: If any file found is not under inputdata_root.
     """
     try:
         with os.scandir(directory) as entries:
             for entry in entries:
                 try:
                     # Recursively process directories (not following symlinks)
                     if entry.is_dir(follow_symlinks=False):
-                        yield from find_owned_files_scandir(entry.path, user_uid)
+                        yield from find_owned_files_scandir(
+                            entry.path, user_uid, inputdata_root
+                        )
 
                     # Things other than directories are handled separately
-                    elif (entry_path := handle_non_dir(entry, user_uid)) is not None:
+                    elif (
+                        entry_path := handle_non_dir(entry, user_uid, inputdata_root)
+                    ) is not None:
                         yield entry_path
 
                 except (OSError, PermissionError) as e:
@@ -119,7 +143,9 @@ def find_owned_files_scandir(directory, user_uid):
         logger.debug("Error accessing %s: %s. Skipping.", directory, e)
 
 
-def replace_files_with_symlinks(source_dir, target_dir, username, dry_run=False):
+def replace_files_with_symlinks(
+    source_dir, target_dir, username, inputdata_root=DEFAULT_SOURCE_ROOT, dry_run=False
+):
     """
     Finds files owned by a specific user in a source directory tree,
     deletes them, and replaces them with symbolic links to the same
@@ -128,6 +154,7 @@ def replace_files_with_symlinks(source_dir, target_dir, username, dry_run=False)
     Args:
         source_dir (str): The root of the directory tree to search for files.
         target_dir (str): The root of the directory tree containing the new files.
+        inputdata_root (str): The root of the directory tree containing CESM input data.
         username (str): The name of the user whose files will be processed.
         dry_run (bool): If True, only show what would be done without making changes.
     """
@@ -152,7 +179,7 @@ def replace_files_with_symlinks(source_dir, target_dir, username, dry_run=False)
     )
 
     # Use efficient scandir-based search
-    for file_path in find_owned_files_scandir(source_dir, user_uid):
+    for file_path in find_owned_files_scandir(source_dir, user_uid, inputdata_root):
         logger.info("Found owned file: %s", file_path)
 
         # Determine the relative path and the new link's destination
@@ -251,6 +278,16 @@ def parse_arguments():
         ),
     )
 
+    # The root of the directory tree containing CESM input data.
+    # ONLY INTENDED FOR USE IN TESTING
+    parser.add_argument(
+        "--inputdata-root",
+        "-inputdata",  # to match rimport
+        type=validate_directory,
+        default=DEFAULT_SOURCE_ROOT,
+        help=argparse.SUPPRESS,
+    )
+
     # Verbosity options (mutually exclusive)
     verbosity_group = parser.add_mutually_exclusive_group()
     verbosity_group.add_argument(
@@ -311,7 +348,11 @@ def main():
 
     # --- Execution ---
     replace_files_with_symlinks(
-        args.source_root, args.target_root, my_username, dry_run=args.dry_run
+        args.source_root,
+        args.target_root,
+        my_username,
+        inputdata_root=args.inputdata_root,
+        dry_run=args.dry_run,
     )
 
     if args.timing:

tests/relink/conftest.py

@@ -7,6 +7,7 @@
 import shutil
 
 import pytest
+from unittest.mock import patch
 
 
 @pytest.fixture(scope="function", name="temp_dirs")
@@ -15,7 +16,8 @@ def fixture_temp_dirs():
     source_dir = tempfile.mkdtemp(prefix="test_source_")
     target_dir = tempfile.mkdtemp(prefix="test_target_")
 
-    yield source_dir, target_dir
+    with patch("relink.DEFAULT_SOURCE_ROOT", source_dir):
+        yield source_dir, target_dir
 
     # Cleanup
     shutil.rmtree(source_dir, ignore_errors=True)

tests/relink/test_args.py

@@ -21,10 +21,11 @@
 
 
 @pytest.fixture(scope="function", name="mock_default_dirs")
-def fixture_mock_default_dirs():
+def fixture_mock_default_dirs(temp_dirs):
     """Mock the default directories to use temporary directories."""
     source_dir = tempfile.mkdtemp(prefix="test_default_source_")
     target_dir = tempfile.mkdtemp(prefix="test_default_target_")
+    source_dir, target_dir = temp_dirs
 
     with patch.object(relink, "DEFAULT_SOURCE_ROOT", source_dir):
         with patch.object(relink, "DEFAULT_TARGET_ROOT", target_dir):

tests/relink/test_cmdline.py

@@ -44,6 +44,8 @@ def test_command_line_execution_dry_run(mock_dirs):
         "--target-root",
         str(target_dir),
         "--dry-run",
+        "--inputdata-root",
+        str(source_dir),
     ]
 
     # Execute the command
@@ -78,6 +80,8 @@ def test_command_line_execution_actual_run(mock_dirs):
         str(source_dir),
         "--target-root",
         str(target_dir),
+        "-inputdata",
+        str(source_dir),
     ]
 
     # Execute the command

tests/relink/test_dryrun.py

@@ -5,6 +5,7 @@
 import os
 import sys
 import logging
+from unittest.mock import patch
 
 import pytest
 
@@ -46,7 +47,7 @@ def test_dry_run_no_changes(dry_run_setup, caplog):
     # Run in dry-run mode
     with caplog.at_level(logging.INFO):
         relink.replace_files_with_symlinks(
-            source_dir, target_dir, username, dry_run=True
+            source_dir, target_dir, username, inputdata_root=source_dir, dry_run=True
         )
 
     # Verify no changes were made
@@ -64,7 +65,7 @@ def test_dry_run_shows_message(dry_run_setup, caplog):
     # Run in dry-run mode
     with caplog.at_level(logging.INFO):
         relink.replace_files_with_symlinks(
-            source_dir, target_dir, username, dry_run=True
+            source_dir, target_dir, username, inputdata_root=source_dir, dry_run=True
         )
 
     # Check that dry-run messages were logged
@@ -80,7 +81,7 @@ def test_dry_run_no_delete_or_create_messages(dry_run_setup, caplog):
     # Run in dry-run mode
     with caplog.at_level(logging.INFO):
         relink.replace_files_with_symlinks(
-            source_dir, target_dir, username, dry_run=True
+            source_dir, target_dir, username, inputdata_root=source_dir, dry_run=True
         )
 
     # Verify actual operation messages are NOT logged

tests/relink/test_find_owned_files_scandir.py

@@ -103,7 +103,9 @@ def test_find_owned_files_basic(temp_dirs):
         f.write("content2")
 
     # Find owned files
-    found_files = list(relink.find_owned_files_scandir(source_dir, user_uid))
+    found_files = list(
+        relink.find_owned_files_scandir(source_dir, user_uid, inputdata_root=source_dir)
+    )
 
     # Verify both files were found
     assert len(found_files) == 2
@@ -130,7 +132,9 @@ def test_find_owned_files_nested(temp_dirs):
             fp.write("content")
 
     # Find owned files
-    found_files = list(relink.find_owned_files_scandir(source_dir, user_uid))
+    found_files = list(
+        relink.find_owned_files_scandir(source_dir, user_uid, inputdata_root=source_dir)
+    )
 
     # Verify all files were found
     assert len(found_files) == 3
@@ -156,7 +160,11 @@ def test_skip_symlinks(temp_dirs, caplog):
 
     # Find owned files with logging
     with caplog.at_level(logging.DEBUG):
-        found_files = list(relink.find_owned_files_scandir(source_dir, user_uid))
+        found_files = list(
+            relink.find_owned_files_scandir(
+                source_dir, user_uid, inputdata_root=source_dir
+            )
+        )
 
     # Verify only regular file was found
     assert len(found_files) == 1
@@ -197,7 +205,11 @@ def test_skip_symlinks_owned_by_different_user(temp_dirs, caplog):
 
     with patch("os.scandir", side_effect=mock_scandir):
         with caplog.at_level(logging.INFO):
-            found_files = list(relink.find_owned_files_scandir(source_dir, user_uid))
+            found_files = list(
+                relink.find_owned_files_scandir(
+                    source_dir, user_uid, inputdata_root=source_dir
+                )
+            )
 
     # Verify only regular file was found
     assert len(found_files) == 1
@@ -216,7 +228,9 @@ def test_empty_directory(temp_dirs):
     user_uid = os.stat(source_dir).st_uid
 
     # Find owned files in empty directory
-    found_files = list(relink.find_owned_files_scandir(source_dir, user_uid))
+    found_files = list(
+        relink.find_owned_files_scandir(source_dir, user_uid, inputdata_root=source_dir)
+    )
 
     # Should return empty list
     assert len(found_files) == 0
@@ -245,7 +259,11 @@ def test_permission_error_handling(temp_dirs, caplog):
     try:
         # Find owned files with debug logging
         with caplog.at_level(logging.DEBUG):
-            found_files = list(relink.find_owned_files_scandir(source_dir, user_uid))
+            found_files = list(
+                relink.find_owned_files_scandir(
+                    source_dir, user_uid, inputdata_root=source_dir
+                )
+            )
 
         # Should find the accessible file but skip the inaccessible directory
         assert file1 in found_files
@@ -272,7 +290,9 @@ def test_only_files_not_directories(temp_dirs):
     os.makedirs(subdir)
 
     # Find owned files
-    found_files = list(relink.find_owned_files_scandir(source_dir, user_uid))
+    found_files = list(
+        relink.find_owned_files_scandir(source_dir, user_uid, inputdata_root=source_dir)
+    )
 
     # Should only find the file, not the directory
     assert len(found_files) == 1
@@ -303,7 +323,11 @@ def test_does_not_follow_symlink_directories(temp_dirs):
         os.symlink(external_dir, symlink_dir)
 
         # Find owned files
-        found_files = list(relink.find_owned_files_scandir(source_dir, user_uid))
+        found_files = list(
+            relink.find_owned_files_scandir(
+                source_dir, user_uid, inputdata_root=source_dir
+            )
+        )
 
         # Should find file in real directory but not in symlinked directory
         assert file_in_real in found_files