Refactor files.HashFile facts

Author: ScottKevill

pyinfra/facts/files.py

@@ -5,6 +5,7 @@
 import re
 import stat
 from datetime import datetime
+from typing import List, Tuple
 
 from pyinfra.api.command import QuoteString, make_formatted_string_command
 from pyinfra.api.facts import FactBase
@@ -204,91 +205,57 @@ class Socket(File):
     type = "socket"
 
 
-class Sha1File(FactBase):
-    """
-    Returns a SHA1 hash of a file. Works with both sha1sum and sha1. Returns
-    ``None`` if the file doest not exist.
-    """
+class HashFileFactBase(FactBase):
+    _raw_cmd: str
+    _regexes: Tuple[str, str]
+
+    def __init_subclass__(cls, digits: int, cmds: List[str], **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
 
-    _regexes = [
-        r"^([a-zA-Z0-9]{40})\s+%s$",
-        r"^SHA1\s+\(%s\)\s+=\s+([a-zA-Z0-9]{40})$",
-    ]
+        raw_hash_cmds = ["%s {0} 2> /dev/null" % cmd for cmd in cmds]
+        raw_hash_cmd = " || ".join(raw_hash_cmds)
+        cls._raw_cmd = "test -e {0} && ( %s ) || true" % raw_hash_cmd
+
+        assert cls.__name__.endswith("File")
+        hash_name = cls.__name__[:-4].upper()
+        cls._regexes = (
+            # GNU coreutils style:
+            r"^([a-fA-F0-9]{%d})\s+%%s$" % digits,
+            # BSD style:
+            r"^%s\s+\(%%s\)\s+=\s+([a-fA-F0-9]{%d})$" % (hash_name, digits),
+        )
 
     def command(self, path):
         self.path = path
-        return make_formatted_string_command(
-            (
-                "test -e {0} && ( "
-                "sha1sum {0} 2> /dev/null || shasum {0} 2> /dev/null || sha1 {0} "
-                ") || true"
-            ),
-            QuoteString(path),
-        )
+        return make_formatted_string_command(self._raw_cmd, QuoteString(path))
 
     def process(self, output):
+        output = output[0]
+        escaped_path = re.escape(self.path)
         for regex in self._regexes:
-            regex = regex % re.escape(self.path)
-            matches = re.match(regex, output[0])
+            matches = re.match(regex % escaped_path, output)
             if matches:
                 return matches.group(1)
 
 
-class Sha256File(FactBase):
+class Sha1File(HashFileFactBase, digits=40, cmds=["sha1sum", "shasum", "sha1"]):
     """
-    Returns a SHA256 hash of a file, or ``None`` if the file does not exist.
+    Returns a SHA1 hash of a file. Works with both sha1sum and sha1. Returns
+    ``None`` if the file doest not exist.
     """
 
-    _regexes = [
-        r"^([a-zA-Z0-9]{64})\s+%s$",
-        r"^SHA256\s+\(%s\)\s+=\s+([a-zA-Z0-9]{64})$",
-    ]
-
-    def command(self, path):
-        self.path = path
-        return make_formatted_string_command(
-            (
-                "test -e {0} && ( "
-                "sha256sum {0} 2> /dev/null "
-                "|| shasum -a 256 {0} 2> /dev/null "
-                "|| sha256 {0} "
-                ") || true"
-            ),
-            QuoteString(path),
-        )
 
-    def process(self, output):
-        for regex in self._regexes:
-            regex = regex % re.escape(self.path)
-            matches = re.match(regex, output[0])
-            if matches:
-                return matches.group(1)
+class Sha256File(HashFileFactBase, digits=64, cmds=["sha256sum", "shasum -a 256", "sha256"]):
+    """
+    Returns a SHA256 hash of a file, or ``None`` if the file does not exist.
+    """
 
 
-class Md5File(FactBase):
+class Md5File(HashFileFactBase, digits=32, cmds=["md5sum", "md5"]):
     """
     Returns an MD5 hash of a file, or ``None`` if the file does not exist.
     """
 
-    _regexes = [
-        r"^([a-zA-Z0-9]{32})\s+%s$",
-        r"^MD5\s+\(%s\)\s+=\s+([a-zA-Z0-9]{32})$",
-    ]
-
-    def command(self, path):
-        self.path = path
-        return make_formatted_string_command(
-            "test -e {0} && ( md5sum {0} 2> /dev/null || md5 {0} ) || true",
-            QuoteString(path),
-        )
-
-    def process(self, output):
-        for regex in self._regexes:
-            regex = regex % re.escape(self.path)
-            matches = re.match(regex, output[0])
-            if matches:
-                return matches.group(1)
-
 
 class FindInFile(FactBase):
     """

tests/facts/files.Md5File/file.json

@@ -1,6 +1,6 @@
 {
     "arg": "myfile",
-    "command": "test -e myfile && ( md5sum myfile 2> /dev/null || md5 myfile ) || true",
+    "command": "test -e myfile && ( md5sum myfile 2> /dev/null || md5 myfile 2> /dev/null ) || true",
     "output": [
         "c10ba97d7c9078a006d26b5db01d8ee7 myfile"
     ],

tests/facts/files.Md5File/file_bsd_style.json

@@ -1,6 +1,6 @@
 {
     "arg": "myfile",
-    "command": "test -e myfile && ( md5sum myfile 2> /dev/null || md5 myfile ) || true",
+    "command": "test -e myfile && ( md5sum myfile 2> /dev/null || md5 myfile 2> /dev/null ) || true",
     "output": [
         "MD5 (myfile) = c10ba97d7c9078a006d26b5db01d8ee7"
     ],

tests/facts/files.Md5File/file_needs_quotes.json

@@ -1,6 +1,6 @@
 {
     "arg": "my () file && special_chars.txt",
-    "command": "test -e 'my () file && special_chars.txt' && ( md5sum 'my () file && special_chars.txt' 2> /dev/null || md5 'my () file && special_chars.txt' ) || true",
+    "command": "test -e 'my () file && special_chars.txt' && ( md5sum 'my () file && special_chars.txt' 2> /dev/null || md5 'my () file && special_chars.txt' 2> /dev/null ) || true",
     "output": [
         "c10ba97d7c9078a006d26b5db01d8ee7 my () file && special_chars.txt"
     ],

tests/facts/files.Sha1File/file.json

@@ -1,6 +1,6 @@
 {
     "arg": "myfile",
-    "command": "test -e myfile && ( sha1sum myfile 2> /dev/null || shasum myfile 2> /dev/null || sha1 myfile ) || true",
+    "command": "test -e myfile && ( sha1sum myfile 2> /dev/null || shasum myfile 2> /dev/null || sha1 myfile 2> /dev/null ) || true",
     "output": [
         "85746ef87ddabd9fdf4836c5835e34a030d2a141 myfile"
     ],

tests/facts/files.Sha1File/file_bsd_style.json

@@ -1,6 +1,6 @@
 {
     "arg": "myfile",
-    "command": "test -e myfile && ( sha1sum myfile 2> /dev/null || shasum myfile 2> /dev/null || sha1 myfile ) || true",
+    "command": "test -e myfile && ( sha1sum myfile 2> /dev/null || shasum myfile 2> /dev/null || sha1 myfile 2> /dev/null ) || true",
     "output": [
         "SHA1 (myfile) = 85746ef87ddabd9fdf4836c5835e34a030d2a141"
     ],

tests/facts/files.Sha1File/file_needs_quotes.json

@@ -1,6 +1,6 @@
 {
     "arg": "my () file && special_chars.txt",
-    "command": "test -e 'my () file && special_chars.txt' && ( sha1sum 'my () file && special_chars.txt' 2> /dev/null || shasum 'my () file && special_chars.txt' 2> /dev/null || sha1 'my () file && special_chars.txt' ) || true",
+    "command": "test -e 'my () file && special_chars.txt' && ( sha1sum 'my () file && special_chars.txt' 2> /dev/null || shasum 'my () file && special_chars.txt' 2> /dev/null || sha1 'my () file && special_chars.txt' 2> /dev/null ) || true",
     "output": [
         "85746ef87ddabd9fdf4836c5835e34a030d2a141 my () file && special_chars.txt"
     ],

tests/facts/files.Sha256File/file.json

@@ -1,6 +1,6 @@
 {
     "arg": "myfile",
-    "command": "test -e myfile && ( sha256sum myfile 2> /dev/null || shasum -a 256 myfile 2> /dev/null || sha256 myfile ) || true",
+    "command": "test -e myfile && ( sha256sum myfile 2> /dev/null || shasum -a 256 myfile 2> /dev/null || sha256 myfile 2> /dev/null ) || true",
     "output": [
         "3867882e8ccc16bd6a1e3e214a46608f9cf5d21687bbb8c25751da3c47b48033 myfile"
     ],

tests/facts/files.Sha256File/file_bsd_style.json

@@ -1,6 +1,6 @@
 {
     "arg": "myfile",
-    "command": "test -e myfile && ( sha256sum myfile 2> /dev/null || shasum -a 256 myfile 2> /dev/null || sha256 myfile ) || true",
+    "command": "test -e myfile && ( sha256sum myfile 2> /dev/null || shasum -a 256 myfile 2> /dev/null || sha256 myfile 2> /dev/null ) || true",
     "output": [
         "SHA256 (myfile) = 3867882e8ccc16bd6a1e3e214a46608f9cf5d21687bbb8c25751da3c47b48033"
     ],

tests/facts/files.Sha256File/file_needs_quotes.json

@@ -1,6 +1,6 @@
 {
     "arg": "my () file && special_chars.txt",
-    "command": "test -e 'my () file && special_chars.txt' && ( sha256sum 'my () file && special_chars.txt' 2> /dev/null || shasum -a 256 'my () file && special_chars.txt' 2> /dev/null || sha256 'my () file && special_chars.txt' ) || true",
+    "command": "test -e 'my () file && special_chars.txt' && ( sha256sum 'my () file && special_chars.txt' 2> /dev/null || shasum -a 256 'my () file && special_chars.txt' 2> /dev/null || sha256 'my () file && special_chars.txt' 2> /dev/null ) || true",
     "output": [
         "3867882e8ccc16bd6a1e3e214a46608f9cf5d21687bbb8c25751da3c47b48033 my () file && special_chars.txt"
     ],