files.Block: fix multiple bugs, added try_prevent_shell_expansion…

… and support for `content: list[str]` (#1030) * #1028 Removed dollar signs, looks like they were typo * #984 Renamed files.marked_block in documentation to files.block * #1031 Adjusted awk to prevent duplications * #1029 Added try_prevent_shell_expansion parameter to partially help with mentioned issue * And also changed type for content to str | list[str]. This makes it more unified with facts.files.Block, that already returns list[str] * Use `OperationValueError` in `files.block` * Update `files.block` tests with fixes --------- Co-authored-by: Nick Mills-Barrett <[email protected]>
pyinfra-dev · Jan 13, 2024 · 6972e3d · 6972e3d
1 parent f0297cd
commit 6972e3d
Show file tree

Hide file tree

Showing 10 changed files with 61 additions and 25 deletions.
diff --git a/pyinfra/operations/files.py b/pyinfra/operations/files.py
@@ -19,6 +19,7 @@
     FileUploadCommand,
     OperationError,
     OperationTypeError,
+    OperationValueError,
     QuoteString,
     RsyncCommand,
     StringCommand,
@@ -1623,6 +1624,7 @@ def block(
     line=None,
     backup=False,
     escape_regex_characters=False,
+    try_prevent_shell_expansion=False,
     before=False,
     after=False,
     marker=None,
@@ -1640,6 +1642,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``
@@ -1656,12 +1659,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",
@@ -1670,7 +1676,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",
@@ -1679,21 +1685,29 @@ 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>",
             regex=".*<body>.*",
             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!")
@@ -1732,22 +1746,36 @@ 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'
             redirect = ">" if (current is None) else ">>"
             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)
             regex = adjust_regex(line, escape_regex_characters)
             print_before = "{ print }" if before else ""
@@ -1757,21 +1785,29 @@ 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,
                 )
@@ -1783,7 +1819,7 @@ def block(
             host.create_fact(
                 Block,
                 kwargs={"path": path, "marker": marker, "begin": begin, "end": end},
-                data=content.split("\n"),
+                data=content,
             )
     else:  # remove the marked_block
         if content:

diff --git a/tests/operations/files.block/add_existing_block_different_content.json b/tests/operations/files.block/add_existing_block_different_content.json
@@ -12,6 +12,6 @@
         }
     },
     "commands": [
-        "OUT=\"$(TMPDIR=/tmp mktemp -t pyinfra.XXXXXX)\" &&  awk 'BEGIN {{f=1; x=ARGV[2]; ARGV[2]=\"\"}}/# BEGIN PYINFRA BLOCK/ {print; print x; f=0} /# END PYINFRA BLOCK/ {print; f=1} f' /home/someone/something $\"should be this\" > $OUT && chmod $(stat -c %a /home/someone/something 2>/dev/null || stat -f %Lp /home/someone/something ) $OUT &&  (chown $(stat -c \"%U:%G\" /home/someone/something 2>/dev/null) $OUT ||  chown -n $(stat -f \"%u:%g\" /home/someone/something ) $OUT)  && mv \"$OUT\" /home/someone/something"
+        "OUT=\"$(TMPDIR=/tmp mktemp -t pyinfra.XXXXXX)\" &&  awk 'BEGIN {{f=1; x=ARGV[2]; ARGV[2]=\"\"}}/# BEGIN PYINFRA BLOCK/ {print; print x; f=0} /# END PYINFRA BLOCK/ {print; f=1; next} f' /home/someone/something \"should be this\" > $OUT && chmod $(stat -c %a /home/someone/something 2>/dev/null || stat -f %Lp /home/someone/something ) $OUT &&  (chown $(stat -c \"%U:%G\" /home/someone/something 2>/dev/null) $OUT ||  chown -n $(stat -f \"%u:%g\" /home/someone/something ) $OUT)  && mv \"$OUT\" /home/someone/something"
     ]
 }
diff --git a/tests/operations/files.block/add_existing_block_different_content_and_backup.json b/tests/operations/files.block/add_existing_block_different_content_and_backup.json
@@ -13,6 +13,6 @@
         }
     },
     "commands": [
-        "cp /home/someone/something /home/someone/something.a-timestamp && OUT=\"$(TMPDIR=/tmp mktemp -t pyinfra.XXXXXX)\" &&  awk 'BEGIN {{f=1; x=ARGV[2]; ARGV[2]=\"\"}}/# BEGIN PYINFRA BLOCK/ {print; print x; f=0} /# END PYINFRA BLOCK/ {print; f=1} f' /home/someone/something $\"should be this\" > $OUT && chmod $(stat -c %a /home/someone/something 2>/dev/null || stat -f %Lp /home/someone/something ) $OUT &&  (chown $(stat -c \"%U:%G\" /home/someone/something 2>/dev/null) $OUT ||  chown -n $(stat -f \"%u:%g\" /home/someone/something ) $OUT)  && mv \"$OUT\" /home/someone/something"
+        "cp /home/someone/something /home/someone/something.a-timestamp && OUT=\"$(TMPDIR=/tmp mktemp -t pyinfra.XXXXXX)\" &&  awk 'BEGIN {{f=1; x=ARGV[2]; ARGV[2]=\"\"}}/# BEGIN PYINFRA BLOCK/ {print; print x; f=0} /# END PYINFRA BLOCK/ {print; f=1; next} f' /home/someone/something \"should be this\" > $OUT && chmod $(stat -c %a /home/someone/something 2>/dev/null || stat -f %Lp /home/someone/something ) $OUT &&  (chown $(stat -c \"%U:%G\" /home/someone/something 2>/dev/null) $OUT ||  chown -n $(stat -f \"%u:%g\" /home/someone/something ) $OUT)  && mv \"$OUT\" /home/someone/something"
     ]
 }
diff --git a/tests/operations/files.block/add_line_given_but_before_eq_after.json b/tests/operations/files.block/add_line_given_but_before_eq_after.json
@@ -13,7 +13,7 @@
         }
     },
     "exception": {
-        "names": ["ValueError"],
+        "names": ["OperationValueError"],
         "message": "only one of 'before' or 'after' used when 'line` is specified"
     }
 }
diff --git a/tests/operations/files.block/add_no_content_provided.json b/tests/operations/files.block/add_no_content_provided.json
@@ -7,7 +7,7 @@
         }
     },
     "exception": {
-        "names": ["ValueError"],
+        "names": ["OperationValueError"],
         "message": "'content' must be supplied when 'present' == True"
     }
 }
diff --git a/tests/operations/files.block/add_no_existing_block_and_no_line.json b/tests/operations/files.block/add_no_existing_block_and_no_line.json
@@ -12,6 +12,6 @@
         }
     },
     "commands": [
-        "cat > /home/someone/something <<PYINFRAHERE\n# BEGIN PYINFRA BLOCK\nplease add this\n# END PYINFRA BLOCK\nPYINFRAHERE"
+        "cat > /home/someone/something <<PYINFRAHERE \n# BEGIN PYINFRA BLOCK\nplease add this\n# END PYINFRA BLOCK\nPYINFRAHERE"
     ]
 }
diff --git a/tests/operations/files.block/add_no_existing_block_line_provided.json b/tests/operations/files.block/add_no_existing_block_line_provided.json
@@ -12,6 +12,6 @@
         }
     },
     "commands": [
-        "OUT=\"$(TMPDIR=/tmp mktemp -t pyinfra.XXXXXX)\" &&  awk 'BEGIN {x=ARGV[2]; ARGV[2]=\"\"}  f!=1 && /^.*before this.*$/ { print x; f=1} END {if (f==0) print ARGV[2] } { print }' /home/someone/something $\"# BEGIN PYINFRA BLOCK\nplease add this\n# END PYINFRA BLOCK\" > $OUT && chmod $(stat -c %a /home/someone/something 2>/dev/null || stat -f %Lp /home/someone/something ) $OUT &&  (chown $(stat -c \"%U:%G\" /home/someone/something 2>/dev/null) $OUT ||  chown -n $(stat -f \"%u:%g\" /home/someone/something ) $OUT)  && mv \"$OUT\" /home/someone/something"
+        "OUT=\"$(TMPDIR=/tmp mktemp -t pyinfra.XXXXXX)\" &&  awk 'BEGIN {x=ARGV[2]; ARGV[2]=\"\"}  f!=1 && /^.*before this.*$/ { print x; f=1} END {if (f==0) print ARGV[2] } { print }' /home/someone/something \"# BEGIN PYINFRA BLOCK\nplease add this\n# END PYINFRA BLOCK\" > $OUT && chmod $(stat -c %a /home/someone/something 2>/dev/null || stat -f %Lp /home/someone/something ) $OUT &&  (chown $(stat -c \"%U:%G\" /home/someone/something 2>/dev/null) $OUT ||  chown -n $(stat -f \"%u:%g\" /home/someone/something ) $OUT)  && mv \"$OUT\" /home/someone/something"
     ]
 }
diff --git a/tests/operations/files.block/add_no_existing_block_line_provided_escape_regex.json b/tests/operations/files.block/add_no_existing_block_line_provided_escape_regex.json
@@ -13,6 +13,6 @@
         }
     },
     "commands": [
-        "OUT=\"$(TMPDIR=/tmp mktemp -t pyinfra.XXXXXX)\" &&  awk 'BEGIN {x=ARGV[2]; ARGV[2]=\"\"}  f!=1 && /^.*before this \\*.*$/ { print x; f=1} END {if (f==0) print ARGV[2] } { print }' /home/someone/something $\"# BEGIN PYINFRA BLOCK\nplease add this\n# END PYINFRA BLOCK\" > $OUT && chmod $(stat -c %a /home/someone/something 2>/dev/null || stat -f %Lp /home/someone/something ) $OUT &&  (chown $(stat -c \"%U:%G\" /home/someone/something 2>/dev/null) $OUT ||  chown -n $(stat -f \"%u:%g\" /home/someone/something ) $OUT)  && mv \"$OUT\" /home/someone/something"
+        "OUT=\"$(TMPDIR=/tmp mktemp -t pyinfra.XXXXXX)\" &&  awk 'BEGIN {x=ARGV[2]; ARGV[2]=\"\"}  f!=1 && /^.*before this \\*.*$/ { print x; f=1} END {if (f==0) print ARGV[2] } { print }' /home/someone/something \"# BEGIN PYINFRA BLOCK\nplease add this\n# END PYINFRA BLOCK\" > $OUT && chmod $(stat -c %a /home/someone/something 2>/dev/null || stat -f %Lp /home/someone/something ) $OUT &&  (chown $(stat -c \"%U:%G\" /home/someone/something 2>/dev/null) $OUT ||  chown -n $(stat -f \"%u:%g\" /home/someone/something ) $OUT)  && mv \"$OUT\" /home/someone/something"
     ]
 }
diff --git a/tests/operations/files.block/add_no_existing_file.json b/tests/operations/files.block/add_no_existing_file.json
@@ -12,6 +12,6 @@
         }
     },
     "commands": [
-        "cat > /home/someone/something <<PYINFRAHERE\n# BEGIN PYINFRA BLOCK\nplease add this\n# END PYINFRA BLOCK\nPYINFRAHERE"
+        "cat > /home/someone/something <<PYINFRAHERE \n# BEGIN PYINFRA BLOCK\nplease add this\n# END PYINFRA BLOCK\nPYINFRAHERE"
     ]
 }
diff --git a/tests/operations/files.block/add_no_line_given_but_before_ne_after.json b/tests/operations/files.block/add_no_line_given_but_before_ne_after.json
@@ -12,7 +12,7 @@
         }
     },
     "exception": {
-        "names": ["ValueError"],
+        "names": ["OperationValueError"],
         "message": "'line' must be supplied or 'before' and 'after' must be equal"
     }
 }