Move generic code to lmi.test.util

1 files changed, 214 insertions, 0 deletions

diff --git a/src/python/lmi/test/util.py b/src/python/lmi/test/util.py
index 241e16a..699507c 100644
--- a/src/python/lmi/test/util.py
+++ b/src/python/lmi/test/util.py
@@ -21,8 +21,14 @@ LMI test utilities.
 """
 
 import os
+import hashlib
 import pywbem
+import random
+import tempfile
 import socket
+import string
+import subprocess
+from collections import OrderedDict
 from lmi.test import unittest
 
 def is_this_system(system_name):
@@ -137,3 +143,211 @@ def check_inames_equal(fst, snd):
 
     return True
 
+def random_string(strength=6, chars=None, prefix=""):
+    """
+    Generate a random string, e.g. usable as UID/GID
+
+    strength is count of random characters in the final string.  chars
+    is sequence of characters to choose from, and prefix can be provided
+    to prepend it to final string.
+    """
+    if chars is None:
+        chars = string.ascii_uppercase + string.digits
+    salt = ''.join([random.choice(chars) for x in range(strength)])
+    return prefix + salt
+
+
+class BackupStorage():
+    """
+    Simple file backup storage.
+
+    * Only supports files.
+    * Only supports absolute paths.
+    * Consecutive backups rewrite each other.
+    * Does not autodestroy the backup.
+    """
+
+    def __init__(self):
+        self.root = tempfile.mkdtemp(prefix=self.__class__.__name__ + ".")
+        self.backups = OrderedDict()
+        subprocess.check_call(["mkdir", "-p", self.root])
+
+    def _copy(self, src, dest):
+        """
+        Copy src to dst --- force, keep meta, no questions asked
+        """
+        subprocess.check_call(["cp", "-a", "-f", src, dest])
+
+    def _get_bpath(self, path):
+        """
+        Take original path and return path to backup.
+        """
+        if not path.startswith("/"):
+            raise ValueError("only absolute paths are supported")
+        digest = hashlib.sha1(path).hexdigest()
+        return self.root + "/" + digest
+
+    def _update_index(self):
+        """
+        Create/update an index file to help in case of backup investigation
+
+        For convenience, index file is sorted by real path.
+        """
+        paths = sorted(self.backups.keys())
+        with open(self.root + "/index", "w+") as fh:
+            for path in paths:
+                fh.write("%s %s\n" % (self.backups[path], path))
+
+    def add_files(self, paths):
+        """
+        Add list of tiles to backup storage
+        """
+        for path in paths:
+            self.add_file(path)
+
+    def add_file(self, path):
+        """
+        Add a file to backup storage
+        """
+        bpath = self._get_bpath(path)
+        self._copy(path, bpath)
+        self.backups[path] = bpath
+        self._update_index()
+
+    def restore(self, path):
+        """
+        Restore particular path
+        """
+        try:
+            self._copy(self.backups[path], path)
+        except KeyError:
+            raise ValueError("path not stored: %s" % path)
+
+    def restore_all(self):
+        """
+        Restore all stored paths in same order as they were stored
+        """
+        for key in self.backups.keys():
+            self.restore(key)
+
+    def destroy_backup(self):
+        """
+        Destroy the temporary backup
+        """
+        subprocess.call(["rm", "-rf", self.root])
+
+
+class BaseCrippler:
+    """
+    Helper class for crippling system files.
+
+    To use the class, you need to sub-class it and implement
+    _define_cases method.
+    """
+
+    LINE_LENGTH = 500
+    LINE_COUNT = 50
+    BINARY_LENGTH = 10 * 1024 * 1024
+
+    ## virtual
+    #
+
+    def _define_cases(self):
+        """
+        Define cases per file supported
+
+        This function must return a dict with one set of cases per
+        file: key is path and value is another dict defining cases
+        as pairs of name ([a-zA-Z_]) and content.
+
+        Quick example:
+
+            {
+                '/etc/file1': {
+                    'case1': "some triggering content",
+                    'case2': "some other triggering content",
+                    'case3': "some funny triggering content",
+                },
+                '/etc/file2': {
+                    'case1': "some triggering content",
+                    'case2': "some other triggering content",
+                    'case3': "some funny triggering content",
+                },
+            }
+
+        Note that trailing newline is added automatically to each content
+        string. Also, whether content will be appended or replaced is decided
+        by caller of the BaseCrippler.cripple method.
+        """
+        pass
+
+    ## internal
+    #
+
+    def __init__(self):
+        self.autocases = {
+            'empty': lambda: '',
+            'random_line': self._random_line,
+            'random_lines': self._random_lines,
+            'random_binary': self._random_binary,
+        }
+        self.cases = self._define_cases()
+
+    def _append_to(self, path, content):
+        with open(path, 'a+') as fh:
+            fh.write(content)
+
+    def _clobber(self, path, content):
+        with open(path, 'w+') as fh:
+            fh.write(content)
+
+    def _random_binary(self, size=BINARY_LENGTH):
+        chars = ''.join([chr(i) for i in xrange(256)])
+        return random_string(strength=size, chars=chars)
+
+    def _random_line(self, size=LINE_LENGTH):
+        chars = string.letters + string.punctuation + " \t"
+        return random_string(strength=size, chars=chars) + "\n"
+
+    def _random_lines(self, size=LINE_LENGTH, count=LINE_COUNT):
+        return "".join([self._random_line(size) for i in xrange(count)])
+
+    def _get_content(self, path, case):
+
+        try:
+            content = self.autocases[case]()
+        except KeyError:
+            try:
+                content = self.cases[path][case] + "\n"
+            except KeyError:
+                raise ValueError("unknown case: %s for: %s" % (case, path))
+        return content
+
+    ## public
+    #
+
+    def all_cases_for(self, path):
+        """
+        Return list of cases available for path
+        """
+        return self.cases[path].keys() + self.autocases.keys()
+
+    def all_paths(self):
+        """
+        Return list of paths served by this implementation
+        """
+        return self.cases.keys()
+
+    def cripple(self, path, case, op="replace"):
+        """
+        Cripple file according to selected case.
+
+        op is either "append" or "replace" and means that the content will
+        be appended to the file, otherwise it will replace it.
+        """
+        if op == 'replace':
+            self._clobber(path, self._get_content(path, case))
+        elif op == 'append':
+            self._append_to(path, self._get_content(path, case))
+        else:
+            raise ValueError("unknown op: %s" % op)