Merge commit from fork

Changes for users:

- (BREAKING) `dotenv.set_key` and `dotenv.unset_key` used to follow
  symlinks in some situations. This is no longer the case.  For that
  behavior to be restored in all cases, `follow_symlinks=True` should be
  used.
- (BREAKING) In the CLI, `set` and `unset` used to follow symlinks in
  some situations. This is no longer the case.
- (BREAKING) `dotenv.set_key`, `dotenv.unset_key` and the CLI commands
  `set` and `unset` used to reset the file mode of the modified .env
  file to `0o600` in some situations. This is no longer the case: The
  original mode of the file is now preserved. Is the file needed to be
  created or wasn't a regular file, mode `0o600` is used.

Author: bbc2

src/dotenv/cli.py

@@ -114,7 +114,13 @@ def list_values(ctx: click.Context, output_format: str) -> None:
 @click.argument("key", required=True)
 @click.argument("value", required=True)
 def set_value(ctx: click.Context, key: Any, value: Any) -> None:
-    """Store the given key/value."""
+    """
+    Store the given key/value.
+
+    This doesn't follow symlinks, to avoid accidentally modifying a file at a
+    potentially untrusted path.
+    """
+
     file = ctx.obj["FILE"]
     quote = ctx.obj["QUOTE"]
     export = ctx.obj["EXPORT"]
@@ -146,7 +152,12 @@ def get(ctx: click.Context, key: Any) -> None:
 @click.pass_context
 @click.argument("key", required=True)
 def unset(ctx: click.Context, key: Any) -> None:
-    """Removes the given key."""
+    """
+    Removes the given key.
+
+    This doesn't follow symlinks, to avoid accidentally modifying a file at a
+    potentially untrusted path.
+    """
     file = ctx.obj["FILE"]
     quote = ctx.obj["QUOTE"]
     success, key = unset_key(file, key, quote)

src/dotenv/main.py

@@ -2,7 +2,6 @@
 import logging
 import os
 import pathlib
-import shutil
 import stat
 import sys
 import tempfile
@@ -14,9 +13,7 @@
 from .variables import parse_variables
 
 # A type alias for a string path to be used for the paths in this file.
-# These paths may flow to `open()` and `shutil.move()`; `shutil.move()`
-# only accepts string paths, not byte paths or file descriptors. See
-# https://github.com/python/typeshed/pull/6832.
+# These paths may flow to `open()` and `os.replace()`.
 StrPath = Union[str, "os.PathLike[str]"]
 
 logger = logging.getLogger(__name__)
@@ -142,21 +139,54 @@ def get_key(
 def rewrite(
     path: StrPath,
     encoding: Optional[str],
+    follow_symlinks: bool = False,
 ) -> Iterator[Tuple[IO[str], IO[str]]]:
-    pathlib.Path(path).touch()
+    if follow_symlinks:
+        path = os.path.realpath(path)
 
-    with tempfile.NamedTemporaryFile(mode="w", encoding=encoding, delete=False) as dest:
+    try:
+        source: IO[str] = open(path, encoding=encoding)
+        try:
+            path_stat = os.lstat(path)
+            original_mode: Optional[int] = (
+                stat.S_IMODE(path_stat.st_mode)
+                if stat.S_ISREG(path_stat.st_mode)
+                else None
+            )
+        except BaseException:
+            source.close()
+            raise
+    except FileNotFoundError:
+        source = io.StringIO("")
+        original_mode = None
+
+    with tempfile.NamedTemporaryFile(
+        mode="w",
+        encoding=encoding,
+        delete=False,
+        prefix=".tmp_",
+        dir=os.path.dirname(os.path.abspath(path)),
+    ) as dest:
+        dest_path = pathlib.Path(dest.name)
         error = None
+
         try:
-            with open(path, encoding=encoding) as source:
+            with source:
                 yield (source, dest)
         except BaseException as err:
             error = err
 
     if error is None:
-        shutil.move(dest.name, path)
+        try:
+            if original_mode is not None:
+                os.chmod(dest_path, original_mode)
+
+            os.replace(dest_path, path)
+        except BaseException:
+            dest_path.unlink(missing_ok=True)
+            raise
     else:
-        os.unlink(dest.name)
+        dest_path.unlink(missing_ok=True)
         raise error from None
 
 
@@ -167,12 +197,16 @@ def set_key(
     quote_mode: str = "always",
     export: bool = False,
     encoding: Optional[str] = "utf-8",
+    follow_symlinks: bool = False,
 ) -> Tuple[Optional[bool], str, str]:
     """
     Adds or Updates a key/value to the given .env
 
-    If the .env path given doesn't exist, fails instead of risking creating
-    an orphan .env somewhere in the filesystem
+    The target .env file is created if it doesn't exist.
+
+    This function doesn't follow symlinks by default, to avoid accidentally
+    modifying a file at a potentially untrusted path. If you don't need this
+    protection and need symlinks to be followed, use `follow_symlinks`.
     """
     if quote_mode not in ("always", "auto", "never"):
         raise ValueError(f"Unknown quote_mode: {quote_mode}")
@@ -190,7 +224,10 @@ def set_key(
     else:
         line_out = f"{key_to_set}={value_out}\n"
 
-    with rewrite(dotenv_path, encoding=encoding) as (source, dest):
+    with rewrite(dotenv_path, encoding=encoding, follow_symlinks=follow_symlinks) as (
+        source,
+        dest,
+    ):
         replaced = False
         missing_newline = False
         for mapping in with_warn_for_invalid_lines(parse_stream(source)):
@@ -213,19 +250,27 @@ def unset_key(
     key_to_unset: str,
     quote_mode: str = "always",
     encoding: Optional[str] = "utf-8",
+    follow_symlinks: bool = False,
 ) -> Tuple[Optional[bool], str]:
     """
     Removes a given key from the given `.env` file.
 
     If the .env path given doesn't exist, fails.
     If the given key doesn't exist in the .env, fails.
+
+    This function doesn't follow symlinks by default, to avoid accidentally
+    modifying a file at a potentially untrusted path. If you don't need this
+    protection and need symlinks to be followed, use `follow_symlinks`.
     """
     if not os.path.exists(dotenv_path):
         logger.warning("Can't delete from %s - it doesn't exist.", dotenv_path)
         return None, key_to_unset
 
     removed = False
-    with rewrite(dotenv_path, encoding=encoding) as (source, dest):
+    with rewrite(dotenv_path, encoding=encoding, follow_symlinks=follow_symlinks) as (
+        source,
+        dest,
+    ):
         for mapping in with_warn_for_invalid_lines(parse_stream(source)):
             if mapping.key == key_to_unset:
                 removed = True

tests/test_main.py

@@ -62,6 +62,86 @@ def test_set_key_encoding(dotenv_path):
     assert dotenv_path.read_text(encoding=encoding) == "a='é'\n"
 
 
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="file mode bits behave differently on Windows"
+)
+def test_set_key_preserves_file_mode(dotenv_path):
+    dotenv_path.write_text("a=x\n")
+    dotenv_path.chmod(0o640)
+    mode_before = stat.S_IMODE(dotenv_path.stat().st_mode)
+
+    dotenv.set_key(dotenv_path, "a", "y")
+
+    mode_after = stat.S_IMODE(dotenv_path.stat().st_mode)
+    assert mode_before == mode_after
+
+
+def test_rewrite_closes_file_handle_on_lstat_failure(tmp_path):
+    dotenv_path = tmp_path / ".env"
+    dotenv_path.write_text("a=x\n")
+    real_open = open
+    opened_handles = []
+
+    def tracking_open(*args, **kwargs):
+        handle = real_open(*args, **kwargs)
+        opened_handles.append(handle)
+        return handle
+
+    with mock.patch("dotenv.main.os.lstat", side_effect=FileNotFoundError):
+        with mock.patch("dotenv.main.open", side_effect=tracking_open):
+            dotenv.set_key(dotenv_path, "a", "x")
+
+    assert opened_handles, "expected at least one file to be opened"
+    assert all(handle.closed for handle in opened_handles)
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="symlinks require elevated privileges on Windows"
+)
+def test_set_key_symlink_to_existing_file(tmp_path):
+    target = tmp_path / "target.env"
+    target.write_text("a=x\n")
+    symlink = tmp_path / ".env"
+    symlink.symlink_to(target)
+
+    dotenv.set_key(symlink, "a", "y")
+
+    assert target.read_text() == "a=x\n"
+    assert not symlink.is_symlink()
+    assert "a='y'" in symlink.read_text()
+    assert stat.S_IMODE(symlink.stat().st_mode) == 0o600
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="symlinks require elevated privileges on Windows"
+)
+def test_set_key_symlink_to_missing_file(tmp_path):
+    target = tmp_path / "nx"
+    symlink = tmp_path / ".env"
+    symlink.symlink_to(target)
+
+    dotenv.set_key(symlink, "a", "x")
+
+    assert not target.exists()
+    assert not symlink.is_symlink()
+    assert symlink.read_text() == "a='x'\n"
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="symlinks require elevated privileges on Windows"
+)
+def test_set_key_follow_symlinks(tmp_path):
+    target = tmp_path / "target.env"
+    target.write_text("a=x\n")
+    symlink = tmp_path / ".env"
+    symlink.symlink_to(target)
+
+    dotenv.set_key(symlink, "a", "y", follow_symlinks=True)
+
+    assert target.read_text() == "a='y'\n"
+    assert symlink.is_symlink()
+
+
 @pytest.mark.skipif(
     sys.platform != "win32" and os.geteuid() == 0,
     reason="Root user can access files even with 000 permissions.",
@@ -195,6 +275,54 @@ def test_unset_non_existent_file(tmp_path):
     )
 
 
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="symlinks require elevated privileges on Windows"
+)
+def test_unset_key_symlink_to_existing_file(tmp_path):
+    target = tmp_path / "target.env"
+    target.write_text("a=x\n")
+    symlink = tmp_path / ".env"
+    symlink.symlink_to(target)
+
+    dotenv.unset_key(symlink, "a")
+
+    assert target.read_text() == "a=x\n"
+    assert not symlink.is_symlink()
+    assert symlink.read_text() == ""
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="symlinks require elevated privileges on Windows"
+)
+def test_unset_key_symlink_to_missing_file(tmp_path):
+    target = tmp_path / "nx"
+    symlink = tmp_path / ".env"
+    symlink.symlink_to(target)
+    logger = logging.getLogger("dotenv.main")
+
+    with mock.patch.object(logger, "warning") as mock_warning:
+        result = dotenv.unset_key(symlink, "a")
+
+    assert result == (None, "a")
+    assert symlink.is_symlink()
+    mock_warning.assert_called_once()
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="symlinks require elevated privileges on Windows"
+)
+def test_unset_key_follow_symlinks(tmp_path):
+    target = tmp_path / "target.env"
+    target.write_text("a=b\n")
+    symlink = tmp_path / ".env"
+    symlink.symlink_to(target)
+
+    dotenv.unset_key(symlink, "a", follow_symlinks=True)
+
+    assert target.read_text() == ""
+    assert symlink.is_symlink()
+
+
 def prepare_file_hierarchy(path):
     """
     Create a temporary folder structure like the following: