pyinfra 3.0.dev0__py2.py3-none-any.whl → 3.0.1__py2.py3-none-any.whl

pyinfra/operations/files.py

@@ -11,17 +11,18 @@ import traceback
 from datetime import timedelta
 from fnmatch import fnmatch
 from io import StringIO
-from typing import Union
+from pathlib import Path
+from typing import IO, Any, Union
 
 from jinja2 import TemplateRuntimeError, TemplateSyntaxError, UndefinedError
 
-import pyinfra
 from pyinfra import host, logger, state
 from pyinfra.api import (
     FileDownloadCommand,
     FileUploadCommand,
     OperationError,
     OperationTypeError,
+    OperationValueError,
     QuoteString,
     RsyncCommand,
     StringCommand,
@@ -30,6 +31,7 @@ from pyinfra.api import (
 from pyinfra.api.command import make_formatted_string_command
 from pyinfra.api.util import (
     get_call_location,
+    get_file_io,
     get_file_sha1,
     get_path_permissions_mode,
     get_template,
@@ -56,22 +58,21 @@ from .util import files as file_utils
 from .util.files import adjust_regex, ensure_mode_int, get_timestamp, sed_replace, unix_path_join
 
 
-@operation(
-    pipeline_facts={"file": "dest"},
-)
+@operation()
 def download(
-    src,
-    dest,
-    user=None,
-    group=None,
-    mode=None,
-    cache_time=None,
+    src: str,
+    dest: str,
+    user: str | None = None,
+    group: str | None = None,
+    mode: str | None = None,
+    cache_time: int | None = None,
     force=False,
-    sha256sum=None,
-    sha1sum=None,
-    md5sum=None,
-    headers=None,
+    sha256sum: str | None = None,
+    sha1sum: str | None = None,
+    md5sum: str | None = None,
+    headers: dict[str, str] | None = None,
     insecure=False,
+    proxy: str | None = None,
 ):
     """
     Download files from remote locations using ``curl`` or ``wget``.
@@ -88,6 +89,7 @@ def download(
     + md5sum: md5 hash to checksum the downloaded file against
     + headers: optional dictionary of headers to set for the HTTP request
     + insecure: disable SSL verification for the HTTP request
+    + proxy: simple HTTP proxy through which we can download files, form `http://<yourproxy>:<port>`
 
     **Example:**
 
@@ -121,8 +123,8 @@ def download(
         if cache_time:
             # Time on files is not tz-aware, and will be the same tz as the server's time,
             # so we can safely remove the tzinfo from the Date fact before comparison.
-            cache_time = host.get_fact(Date).replace(tzinfo=None) - timedelta(seconds=cache_time)
-            if info["mtime"] and info["mtime"] < cache_time:
+            ctime = host.get_fact(Date).replace(tzinfo=None) - timedelta(seconds=cache_time)
+            if info["mtime"] and info["mtime"] < ctime:
                 download = True
 
         if sha1sum:
@@ -139,10 +141,16 @@ def download(
 
     # If we download, always do user/group/mode as SSH user may be different
     if download:
-        temp_file = state.get_temp_filename(dest)
+        temp_file = host.get_temp_filename(dest)
 
         curl_args: list[Union[str, StringCommand]] = ["-sSLf"]
         wget_args: list[Union[str, StringCommand]] = ["-q"]
+
+        if proxy:
+            curl_args.append(f"--proxy {proxy}")
+            wget_args.append("-e use_proxy=yes")
+            wget_args.append(f"-e http_proxy={proxy}")
+
         if insecure:
             curl_args.append("--insecure")
             wget_args.append("--no-check-certificate")
@@ -219,11 +227,11 @@ def download(
 
 @operation()
 def line(
-    path,
-    line,
+    path: str,
+    line: str,
     present=True,
-    replace=None,
-    flags=None,
+    replace: str | None = None,
+    flags: list[str] | None = None,
     backup=False,
     interpolate_variables=False,
     escape_regex_characters=False,
@@ -261,7 +269,7 @@ def line(
         it will be append to the end of the file.
 
     Ensure new line:
-        This will ensure that the ``line`` being appended is always on a seperate new
+        This will ensure that the ``line`` being appended is always on a separate new
         line in case the file doesn't end with a newline character.
 
 
@@ -375,48 +383,27 @@ def line(
 
     # No line and we want it, append it
     if not present_lines and present:
-        # If the file does not exist - it *might* be created, so we handle it
-        # dynamically with a little script.
-        if present_lines is None:
-            yield make_formatted_string_command(
-                """
-                    if [ -f '{target}' ]; then
-                        ( grep {match_line} '{target}' && \
-                        {sed_replace_command}) 2> /dev/null || \
-                        {echo_command} ;
-                    else
-                        {echo_command} ;
-                    fi
-                """,
-                target=QuoteString(path),
-                match_line=QuoteString(match_line),
-                echo_command=echo_command,
-                sed_replace_command=sed_replace_command,
+        # If we're doing replacement, only append if the *replacement* line
+        # does not exist (as we are appending the replacement).
+        if replace:
+            # Ensure replace explicitly matches a whole line
+            replace_line = replace
+            if not replace_line.startswith("^"):
+                replace_line = f"^{replace_line}"
+            if not replace_line.endswith("$"):
+                replace_line = f"{replace_line}$"
+
+            present_lines = host.get_fact(
+                FindInFile,
+                path=path,
+                pattern=replace_line,
+                interpolate_variables=interpolate_variables,
             )
 
-        # File exists but has no matching lines - append it.
+        if not present_lines:
+            yield echo_command
         else:
-            # If we're doing replacement, only append if the *replacement* line
-            # does not exist (as we are appending the replacement).
-            if replace:
-                # Ensure replace explicitly matches a whole line
-                replace_line = replace
-                if not replace_line.startswith("^"):
-                    replace_line = f"^{replace_line}"
-                if not replace_line.endswith("$"):
-                    replace_line = f"{replace_line}$"
-
-                present_lines = host.get_fact(
-                    FindInFile,
-                    path=path,
-                    pattern=replace_line,
-                    interpolate_variables=interpolate_variables,
-                )
-
-            if not present_lines:
-                yield echo_command
-            else:
-                host.noop('line "{0}" exists in {1}'.format(replace or line, path))
+            host.noop('line "{0}" exists in {1}'.format(replace or line, path))
 
     # Line(s) exists and we want to remove them, replace with nothing
     elif present_lines and not present:
@@ -440,10 +427,10 @@ def line(
 
 @operation()
 def replace(
-    path,
-    text=None,
-    replace=None,
-    flags=None,
+    path: str,
+    text: str | None = None,
+    replace: str | None = None,
+    flags: list[str] | None = None,
     backup=False,
     interpolate_variables=False,
     match=None,  # deprecated
@@ -512,19 +499,17 @@ def replace(
         host.noop('string "{0}" does not exist in {1}'.format(text, path))
 
 
-@operation(
-    pipeline_facts={"find_files": "destination"},
-)
+@operation()
 def sync(
-    src,
-    dest,
-    user=None,
-    group=None,
-    mode=None,
-    dir_mode=None,
+    src: str,
+    dest: str,
+    user: str | None = None,
+    group: str | None = None,
+    mode: str | None = None,
+    dir_mode: str | None = None,
     delete=False,
-    exclude=None,
-    exclude_dir=None,
+    exclude: str | list[str] | tuple[str] | None = None,
+    exclude_dir: str | list[str] | tuple[str] | None = None,
     add_deploy_dir=True,
 ):
     """
@@ -586,7 +571,7 @@ def sync(
     put_files = []
     ensure_dirnames = []
     for dirpath, dirnames, filenames in os.walk(src, topdown=True):
-        remote_dirpath = os.path.normpath(os.path.relpath(dirpath, src))
+        remote_dirpath = Path(os.path.normpath(os.path.relpath(dirpath, src))).as_posix()
 
         # Filter excluded dirs
         for child_dir in dirnames[:]:
@@ -620,7 +605,7 @@ def sync(
     if dest_link_info:
         dest_to_ensure = dest_link_info["link_target"]
 
-    yield from directory(
+    yield from directory._inner(
         path=dest_to_ensure,
         user=user,
         group=group,
@@ -629,7 +614,7 @@ def sync(
 
     # Ensure any remote dirnames
     for dir_path_curr, dir_mode_curr in ensure_dirnames:
-        yield from directory(
+        yield from directory._inner(
             path=unix_path_join(dest, dir_path_curr),
             user=user,
             group=group,
@@ -638,7 +623,7 @@ def sync(
 
     # Put each file combination
     for local_filename, remote_filename in put_files:
-        yield from put(
+        yield from put._inner(
             src=local_filename,
             dest=remote_filename,
             user=user,
@@ -658,7 +643,7 @@ def sync(
             if exclude and any(fnmatch(filename, match) for match in exclude):
                 continue
 
-            yield from file(path=filename, present=False)
+            yield from file._inner(path=filename, present=False)
 
 
 @memoize
@@ -667,7 +652,7 @@ def show_rsync_warning():
 
 
 @operation(is_idempotent=False)
-def rsync(src, dest, flags=["-ax", "--delete"]):
+def rsync(src: str, dest: str, flags: list[str] | None = None):
     """
     Use ``rsync`` to sync a local directory to the remote system. This operation will actually call
     the ``rsync`` binary on your system.
@@ -682,6 +667,8 @@ def rsync(src, dest, flags=["-ax", "--delete"]):
         global arguments.
     """
 
+    if flags is None:
+        flags = ["-ax", "--delete"]
     show_rsync_warning()
 
     try:
@@ -692,11 +679,11 @@ def rsync(src, dest, flags=["-ax", "--delete"]):
     yield RsyncCommand(src, dest, flags)
 
 
-def _create_remote_dir(state, host, remote_filename, user, group):
+def _create_remote_dir(remote_filename, user, group):
     # Always use POSIX style path as local might be Windows, remote always *nix
     remote_dirname = posixpath.dirname(remote_filename)
     if remote_dirname:
-        yield from directory(
+        yield from directory._inner(
             path=remote_dirname,
             user=user,
             group=group,
@@ -709,14 +696,10 @@ def _create_remote_dir(state, host, remote_filename, user, group):
     # We don't (currently) cache the local state, so there's nothing we can
     # update to flag the local file as present.
     is_idempotent=False,
-    pipeline_facts={
-        "file": "src",
-        "sha1_file": "src",
-    },
 )
 def get(
-    src,
-    dest,
+    src: str,
+    dest: str,
     add_deploy_dir=True,
     create_local_dir=False,
     force=False,
@@ -757,11 +740,11 @@ def get(
 
     # No remote file, so assume exists and download it "blind"
     if not remote_file or force:
-        yield FileDownloadCommand(src, dest, remote_temp_filename=state.get_temp_filename(dest))
+        yield FileDownloadCommand(src, dest, remote_temp_filename=host.get_temp_filename(dest))
 
     # No local file, so always download
     elif not os.path.exists(dest):
-        yield FileDownloadCommand(src, dest, remote_temp_filename=state.get_temp_filename(dest))
+        yield FileDownloadCommand(src, dest, remote_temp_filename=host.get_temp_filename(dest))
 
     # Remote file exists - check if it matches our local
     else:
@@ -770,21 +753,16 @@ def get(
 
         # Check sha1sum, upload if needed
         if local_sum != remote_sum:
-            yield FileDownloadCommand(src, dest, remote_temp_filename=state.get_temp_filename(dest))
+            yield FileDownloadCommand(src, dest, remote_temp_filename=host.get_temp_filename(dest))
 
 
-@operation(
-    pipeline_facts={
-        "file": "dest",
-        "sha1_file": "dest",
-    },
-)
+@operation()
 def put(
-    src,
-    dest,
-    user=None,
-    group=None,
-    mode=None,
+    src: str | IO[Any],
+    dest: str,
+    user: str | None = None,
+    group: str | None = None,
+    mode: int | str | bool | None = None,
     add_deploy_dir=True,
     create_remote_dir=True,
     force=False,
@@ -845,6 +823,7 @@ def put(
 
     # Assume string filename
     else:
+        assert isinstance(src, (str, Path))
         # Add deploy directory?
         if add_deploy_dir and state.cwd:
             src = os.path.join(state.cwd, src)
@@ -859,7 +838,7 @@ def put(
             raise IOError("No such file: {0}".format(local_file))
 
     if mode is True:
-        if os.path.isfile(local_file):
+        if isinstance(local_file, str) and os.path.isfile(local_file):
             mode = get_path_permissions_mode(local_file)
         else:
             logger.warning(
@@ -872,19 +851,20 @@ def put(
 
     remote_file = host.get_fact(File, path=dest)
 
-    if not remote_file and host.get_fact(Directory, path=dest):
+    if not remote_file and bool(host.get_fact(Directory, path=dest)):
+        assert isinstance(src, str)
         dest = unix_path_join(dest, os.path.basename(src))
         remote_file = host.get_fact(File, path=dest)
 
     if create_remote_dir:
-        yield from _create_remote_dir(state, host, dest, user, group)
+        yield from _create_remote_dir(dest, user, group)
 
     # No remote file, always upload and user/group/mode if supplied
     if not remote_file or force:
         yield FileUploadCommand(
             local_file,
             dest,
-            remote_temp_filename=state.get_temp_filename(dest),
+            remote_temp_filename=host.get_temp_filename(dest),
         )
 
         if user or group:
@@ -902,7 +882,7 @@ def put(
             yield FileUploadCommand(
                 local_file,
                 dest,
-                remote_temp_filename=state.get_temp_filename(dest),
+                remote_temp_filename=host.get_temp_filename(dest),
             )
 
             if user or group:
@@ -929,7 +909,15 @@ def put(
 
 
 @operation()
-def template(src, dest, user=None, group=None, mode=None, create_remote_dir=True, **data):
+def template(
+    src: str | IO[Any],
+    dest: str,
+    user: str | None = None,
+    group: str | None = None,
+    mode: str | None = None,
+    create_remote_dir=True,
+    **data,
+):
     '''
     Generate a template using jinja2 and write it to the remote system.
 
@@ -1011,7 +999,6 @@ def template(src, dest, user=None, group=None, mode=None, create_remote_dir=True
     data.setdefault("host", host)
     data.setdefault("state", state)
     data.setdefault("inventory", state.inventory)
-    data.setdefault("facts", pyinfra.facts)
 
     # Render and make file-like it's output
     try:
@@ -1026,7 +1013,7 @@ def template(src, dest, user=None, group=None, mode=None, create_remote_dir=True
         line_number = trace_frames[-1][1]
 
         # Quickly read the line in question and one above/below for nicer debugging
-        with open(src, "r") as f:
+        with get_file_io(src, "r") as f:
             template_lines = f.readlines()
 
         template_lines = [line.strip() for line in template_lines]
@@ -1039,14 +1026,14 @@ def template(src, dest, user=None, group=None, mode=None, create_remote_dir=True
                 e,
                 "\n".join(relevant_lines),
             ),
-        )
+        ) from None
 
     output_file = StringIO(output)
     # Set the template attribute for nicer debugging
     output_file.template = src  # type: ignore[attr-defined]
 
     # Pass to the put function
-    yield from put(
+    yield from put._inner(
         src=output_file,
         dest=dest,
         user=user,
@@ -1078,20 +1065,18 @@ def _raise_or_remove_invalid_path(fs_type, path, force, force_backup, force_back
         raise OperationError("{0} exists and is not a {1}".format(path, fs_type))
 
 
-@operation(
-    pipeline_facts={"link": "path"},
-)
+@operation()
 def link(
-    path,
-    target=None,
+    path: str,
+    target: str | None = None,
     present=True,
-    user=None,
-    group=None,
+    user: str | None = None,
+    group: str | None = None,
     symbolic=True,
     create_remote_dir=True,
     force=False,
     force_backup=True,
-    force_backup_dir=None,
+    force_backup_dir: str | None = None,
 ):
     """
     Add/remove/update links.
@@ -1148,7 +1133,6 @@ def link(
     if symbolic:
         add_args.append("-s")
 
-    add_cmd = StringCommand(" ".join(add_args), QuoteString(target), QuoteString(path))
     remove_cmd = StringCommand("rm", "-f", QuoteString(path))
 
     if not present:
@@ -1158,9 +1142,12 @@ def link(
             host.noop("link {link} does not exist")
         return
 
+    assert target is not None  # appease typing QuoteString below
+    add_cmd = StringCommand(" ".join(add_args), QuoteString(target), QuoteString(path))
+
     if info is None:  # create
         if create_remote_dir:
-            yield from _create_remote_dir(state, host, path, user, group)
+            yield from _create_remote_dir(path, user, group)
 
         yield add_cmd
 
@@ -1185,20 +1172,18 @@ def link(
             host.noop("link {0} already exists".format(path))
 
 
-@operation(
-    pipeline_facts={"file": "path"},
-)
+@operation()
 def file(
-    path,
+    path: str,
     present=True,
-    user=None,
-    group=None,
-    mode=None,
+    user: str | None = None,
+    group: str | None = None,
+    mode: int | str | None = None,
     touch=False,
     create_remote_dir=True,
     force=False,
     force_backup=True,
-    force_backup_dir=None,
+    force_backup_dir: str | None = None,
 ):
     """
     Add/remove/update files.
@@ -1259,7 +1244,7 @@ def file(
 
     if info is None:  # create
         if create_remote_dir:
-            yield from _create_remote_dir(state, host, path, user, group)
+            yield from _create_remote_dir(path, user, group)
 
         yield StringCommand("touch", QuoteString(path))
 
@@ -1289,19 +1274,17 @@ def file(
             host.noop("file {0} already exists".format(path))
 
 
-@operation(
-    pipeline_facts={"directory": "path"},
-)
+@operation()
 def directory(
-    path,
+    path: str,
     present=True,
-    user=None,
-    group=None,
-    mode=None,
+    user: str | None = None,
+    group: str | None = None,
+    mode: int | str | None = None,
     recursive=False,
     force=False,
     force_backup=True,
-    force_backup_dir=None,
+    force_backup_dir: str | None = None,
     _no_check_owner_mode=False,
     _no_fail_on_link=False,
 ):
@@ -1393,8 +1376,8 @@ def directory(
             host.noop("directory {0} already exists".format(path))
 
 
-@operation(pipeline_facts={"flags": "path"})
-def flags(path, flags=None, present=True):
+@operation()
+def flags(path: str, flags: list[str] | None = None, present=True):
     """
     Set/clear file flags.
 
@@ -1441,21 +1424,20 @@ def flags(path, flags=None, present=True):
             )
 
 
-@operation(
-    pipeline_facts={"file": "path"},
-)
+@operation()
 def block(
-    path,
-    content=None,
+    path: str,
+    content: str | list[str] | None = None,
     present=True,
-    line=None,
+    line: str | None = None,
     backup=False,
     escape_regex_characters=False,
+    try_prevent_shell_expansion=False,
     before=False,
     after=False,
-    marker=None,
-    begin=None,
-    end=None,
+    marker: str | None = None,
+    begin: str | None = None,
+    end: str | None = None,
 ):
     """
     Ensure content, surrounded by the appropriate markers, is present (or not) in the file.
@@ -1468,6 +1450,7 @@ def block(
     + line: regex before or after which the content should be added if it doesn't exist.
     + backup: whether to backup the file (see ``files.line``). Default False.
     + escape_regex_characters: whether to escape regex characters from the matching line
+    + try_prevent_shell_expansion: tries to prevent shell expanding by values like `$`
     + marker: the base string used to mark the text.  Default is ``# {mark} PYINFRA BLOCK``
     + begin: the value for ``{mark}`` in the marker before the content. Default is ``BEGIN``
     + end: the value for ``{mark}`` in the marker after the content. Default is ``END``
@@ -1484,12 +1467,15 @@ def block(
 
     Removal ignores ``content`` and ``line``
 
+    Preventing shell expansion works by wrapping the content in '`' before passing to `awk`.
+    WARNING: This will break if the content contains raw single quotes.
+
     **Examples:**
 
     .. code:: python
 
         # add entry to /etc/host
-        files.marked_block(
+        files.block(
             name="add IP address for red server",
             path="/etc/hosts",
             content="10.0.0.1 mars-one",
@@ -1498,7 +1484,7 @@ def block(
         )
 
         # have two entries in /etc/host
-        files.marked_block(
+        files.block(
             name="add IP address for red server",
             path="/etc/hosts",
             content="10.0.0.1 mars-one\\n10.0.0.2 mars-two",
@@ -1507,14 +1493,14 @@ def block(
         )
 
         # remove marked entry from /etc/hosts
-        files.marked_block(
+        files.block(
             name="remove all 10.* addresses from /etc/hosts",
             path="/etc/hosts",
             present=False
         )
 
         # add out of date warning to web page
-        files.marked_block(
+        files.block(
             name="add out of date warning to web page",
             path="/var/www/html/something.html",
             content= "<p>Warning: this page is out of date.</p>",
@@ -1522,6 +1508,14 @@ def block(
             after=True
             marker="<!-- {mark} PYINFRA BLOCK -->",
         )
+
+        # put complex alias into .zshrc
+        files.block(
+            path="/home/user/.zshrc",
+            content="eval $(thefuck -a)",
+            try_prevent_shell_expansion=True,
+            marker="## {mark} ALIASES ##"
+        )
     """
 
     logger.warning("The `files.block` operation is currently in beta!")
@@ -1560,14 +1554,23 @@ def block(
     cmd = None
     if present:
         if not content:
-            raise ValueError("'content' must be supplied when 'present' == True")
+            raise OperationValueError("'content' must be supplied when 'present' == True")
         if line:
             if before == after:
-                raise ValueError("only one of 'before' or 'after' used when 'line` is specified")
+                raise OperationValueError(
+                    "only one of 'before' or 'after' used when 'line` is specified"
+                )
         elif before != after:
-            raise ValueError("'line' must be supplied or 'before' and 'after' must be equal")
+            raise OperationValueError(
+                "'line' must be supplied or 'before' and 'after' must be equal"
+            )
+        if isinstance(content, str):
+            # convert string to list of lines
+            content = content.split("\n")
+        if try_prevent_shell_expansion and any("'" in line for line in content):
+            logger.warning("content contains single quotes, shell expansion prevention may fail")
 
-        the_block = "\n".join([mark_1, content, mark_2])
+        the_block = "\n".join([mark_1, *content, mark_2])
 
         if (current is None) or ((current == []) and (before == after)):
             # a) no file or b) file but no markers and we're adding at start or end. Both use 'cat'
@@ -1575,8 +1578,14 @@ def block(
             stdin = "- " if ((current == []) and before) else ""
             # here = hex(random.randint(0, 2147483647))
             here = "PYINFRAHERE"
-            cmd = StringCommand(f"cat {stdin}{redirect}", q_path, f"<<{here}\n{the_block}\n{here}")
+            cmd = StringCommand(
+                f"cat {stdin}{redirect}",
+                q_path,
+                f"<<{here}" if not try_prevent_shell_expansion else f"<<'{here}'",
+                f"\n{the_block}\n{here}",
+            )
         elif current == []:  # markers not found and have a pattern to match (not start or end)
+            assert isinstance(line, str)
             regex = adjust_regex(line, escape_regex_characters)
             print_before = "{ print }" if before else ""
             print_after = "{ print }" if after else ""
@@ -1585,21 +1594,31 @@ def block(
                 f"{print_after} f!=1 && /{regex}/ {{ print x; f=1}} "
                 f"END {{if (f==0) print ARGV[2] }} {print_before}'"
             )
-            cmd = StringCommand(out_prep, prog, q_path, f'$"{the_block}"', "> $OUT &&", real_out)
+            cmd = StringCommand(
+                out_prep,
+                prog,
+                q_path,
+                f'"{the_block}"' if not try_prevent_shell_expansion else f"'{the_block}'",
+                "> $OUT &&",
+                real_out,
+            )
         else:
-            pieces = content.split("\n")
-            if (len(current) != len(pieces)) or (
-                not all(lines[0] == lines[1] for lines in zip(pieces, current))
+            if (len(current) != len(content)) or (
+                not all(lines[0] == lines[1] for lines in zip(content, current))
             ):  # marked_block found but text is different
                 prog = (
                     'awk \'BEGIN {{f=1; x=ARGV[2]; ARGV[2]=""}}'
-                    f"/{mark_1}/ {{print; print x; f=0}} /{mark_2}/ {{print; f=1}} f'"
+                    f"/{mark_1}/ {{print; print x; f=0}} /{mark_2}/ {{print; f=1; next}} f'"
                 )
                 cmd = StringCommand(
                     out_prep,
                     prog,
                     q_path,
-                    '$"' + content + '"',
+                    (
+                        '"' + "\n".join(content) + '"'
+                        if not try_prevent_shell_expansion
+                        else "'" + "\n".join(content) + "'"
+                    ),
                     "> $OUT &&",
                     real_out,
                 )