[Qemu-devel] [RFC] contribs: Add snapshot_tree .py a tool to build snapshot graphs.

[Benoît Canet]

This tool output in dot(1) or D3.js JSON format a forest of qemu images and
snapshot showing their relations given a list of directories.

Signed-off-by: Benoit Canet <address@hidden>
---
 contribs/snapshot_tree.py |  406 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 406 insertions(+)
 create mode 100755 contribs/snapshot_tree.py

diff --git a/contribs/snapshot_tree.py b/contribs/snapshot_tree.py
new file mode 100755
index 0000000..6e3fa50
--- /dev/null
+++ b/contribs/snapshot_tree.py
@@ -0,0 +1,406 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*
+#
+# Recurse in a directory and build dot(1) or d3.js JSON output of the snapshots
+# in it.
+#
+# Copyright (C) 2010 Nodalink, SARL.
+#
+# Author:
+#   Benoît Canet <address@hidden>
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#
+import argparse
+import json
+import os
+import subprocess
+import sys
+import unittest
+
+class Node(object):
+    """ An element representing a snapshot in the tree """
+
+    def __init__(self, path, infos):
+        self._infos = infos
+        self._children = []
+        self._parent = None
+        self._path = path
+
+    def add_child(self, node):
+        """ Add a child to this node """
+        self._children.append(node)
+
+    def set_parent(self, node):
+        """ Return the parent node of this node """
+        self._parent = node
+
+    @property
+    def image(self):
+        """ image accessor """
+        return self._infos['image']
+
+    @property
+    def backing_file(self):
+        """ backing file accessor """
+        return self._infos['backing_file']
+
+    @property
+    def path(self):
+        """ path accessor """
+        return self._path
+
+    @property
+    def children(self):
+        """ children accessor """
+        return self._children
+
+    def jsonable(self):
+        """ Serialize a node in JSON """
+        result = {}
+        if len(self._children):
+            result["children"] = self._children
+        result["name"] = self.image
+        result["size"] = self._infos.get("disk_size", 0)
+        return result
+
+def get_full_path(dirpath, filename):
+    """ Get the full path """
+    return os.path.relpath(os.path.join(dirpath, filename))
+
+def filter_image_list(filenames):
+    """ Remove all non image files """
+    formats = ["cow", "dmg", "parallels", "qcow2", "qcow",
+               "qed", "raw", "vdi", "vmdk", "vpc" ]
+    result = []
+    for filename in filenames:
+        arr = filename.split(".")
+        extension = arr[-1]
+        if extension in formats: 
+            result.append(filename)
+    return result 
+
+class Forest(object):
+    """ A forest of snapshot nodes """
+
+    def __init__(self, directories):
+        self._directories = directories
+        self._node_by_path = {}
+        self._root_node = None
+        self._top_of_forest = {}
+        self._labels = {}
+
+    def jsonable(self):
+        """ Serialize a forest in JSON """
+        return self._root_node
+
+
+    def _build_nodes(self, dirpath, image_list):
+        """ Build all the nodes for a given image list """
+        for filename in image_list:
+            full_path = get_full_path(dirpath, filename)
+            infos, result = get_image_infos(full_path)
+            if not result:
+                perror("Getting image informations")
+                sys.exit(3)
+            # build node 
+            node = Node(full_path, infos)
+            # store in dicts for later usage
+            self._node_by_path[node.path] = node
+            self._top_of_forest[node.path] = node
+
+    def _create_subtrees_and_filter(self):
+        """ Remove the nodes having a parent from the top of the forest """
+        for node in self._node_by_path.itervalues():
+            parent = self._node_by_path.get(node.backing_file, None)
+            if parent:
+                parent.add_child(node)
+                node.set_parent(parent)
+                del self._top_of_forest[node.image]
+
+    def _walk_and_build_forest(self):
+        """ Walk into the given directory and build the forest
+           
+            Work is done in two pass : building the node
+                                       selecting the top at the top
+        
+        """
+        #create the list of images and build nodes
+        for directory in self._directories:
+            cwd = os.getcwd()
+            array = os.path.realpath(directory).split(os.sep)
+            directory = str.join(os.sep, array[:-1])
+            os.chdir(directory)
+            for (dirpath, dirnames, filenames) in os.walk(directory):
+                image_list = filter_image_list(filenames)
+                self._build_nodes(dirpath, image_list)
+            os.chdir(cwd)
+
+        # remove node having a parent from top of the forest
+        self._create_subtrees_and_filter()
+
+        # create root node of the forest
+        self._root_node = Node("", {'image': "/"})
+        for node in self._top_of_forest.itervalues():
+            self._root_node.add_child(node)            
+
+    def _dump_node(self, node):
+        """ Dump a node """
+        self._labels[id(node)] = node.image
+        string = "%i" % (id(node))
+        sys.stdout.write(string)       
+
+    def _dump(self, node):
+        """ Traverse the tree in preorder and dump it """
+        if len(node.children) == 0:
+            self._dump_node(node)
+            sys.stdout.write(";\n")
+        for child in node.children:
+            self._dump_node(node)
+            sys.stdout.write(" -> ")
+            self._dump(child)
+
+    def build(self):
+        """ Build the forest """
+        self._walk_and_build_forest()
+
+    def dump(self, dump_json):
+        """ Dump the forest in dot(1) or d3 JSON representation """
+        # D3 JSON
+        if dump_json:
+            sys.stdout.write(self.to_json())
+            return
+        # dot(1)
+        sys.stdout.write("digraph graphname {\n")
+        node = self._root_node
+        self._dump(node)
+        for key, value in self._labels.iteritems():
+            sys.stdout.write("%i [ label=\"%s\" ];\n" % (key, value))
+        sys.stdout.write("}\n")
+
+    def to_json(self):
+        """ Dump the forest in D3.js JSON format """
+        return json.dumps(self, default=json_complex_handler,
+                          sort_keys=True, indent=2)
+
+def json_complex_handler(obj):
+    """ Handle complex object when serializing to JSON """
+    if hasattr(obj, 'jsonable'):
+        return obj.jsonable()
+    else:
+        raise TypeError, \
+              'object of type %s with value of %s is not JSON serializable' % \
+              (type(obj), repr(obj))
+
+def get_image_infos(file_path):
+    """ Parse the output of qemu-img info 
+
+        qemu-img info should be able to output JSON
+
+    """
+    infos = {'backing_file': None}
+    process = subprocess.Popen(["qemu-img", "info", file_path],
+                          stdout=subprocess.PIPE, close_fds=False)
+
+    line = process.stdout.readline()
+    while line:
+        line = line.strip()
+        line = line.replace(" ", "_")
+        array = line.split(":")
+        if len(array) == 3:
+            key = array[0]
+            value = array[1].split(" ")[0].split("(")[0]
+        else:
+            key = array[0]
+            value = array[1].strip()
+        if key == "virtual_size":
+            value = value.split("(")[0]
+        infos[key] = value.replace("_", '')
+        line = process.stdout.readline()
+
+
+    infos["image"] = os.path.relpath(infos["image"])
+    if infos["backing_file"]:
+        infos["backing_file"] = os.path.relpath(infos["backing_file"])
+
+    infos['virtual_size'] = convert_to_byte(infos['virtual_size'])
+    infos['disk_size'] = convert_to_byte(infos['disk_size'])
+    if infos.has_key("cluster_size"):
+        infos['cluster_size'] = int(infos['cluster_size'])
+
+    process.communicate()
+    if process.returncode != 0:
+        return (None, False)
+
+    return (infos, True)
+
+def is_directory_valid(directory):
+    """ Check that a directory is valid and accessible """
+    if not os.path.exists(directory):
+        return perror("%s does not exists", directory)
+
+    if not os.path.isdir(directory):
+        return perror("%s is not a directory", directory)
+
+    if not os.access(directory, os.R_OK):
+        return perror("Does not have access to %s", directory)
+
+    return True        
+
+def unit_to_byte(unit):
+    """ Convert a unit to byte """
+    units = [ "B", "K", "M", "G", "T", "P"]
+    index = units.index(unit)
+    result = 1
+    for j in range(0, index):
+        result = result * 1024
+    return result
+
+def convert_to_byte(string):
+    """ Convert to byte a human representation """
+    numeric = ""
+    unit = ""
+    for char in string:
+        if char.isdigit() or char == '.':
+            numeric += char
+        if char.isalpha():
+            unit += char
+    numeric = float(numeric)
+    result = 0
+    if numeric:
+        result = numeric * unit_to_byte(unit)
+
+    return int(result)
+
+def perror(*args):
+    """Print an error and return False
+
+    Prefix the error message with Error and end it with . and a carriage
+    return
+
+    """
+    message = "Error: %s.%c"% (args[0], '\n')
+    sys.stderr.write(message % (args[1:]))
+    return False
+
+def run_command(directories, dump_json):
+    """ Function doing the real work """
+    for directory in directories:
+        if not is_directory_valid(directory):
+            sys.exit(1)
+
+    forest = Forest(directories)
+    forest.build()
+    forest.dump(dump_json)
+
+    sys.exit(0)
+
+class Tests(unittest.TestCase):
+    """ Some basic tests to check the code """
+
+    TEST_DIR = "tests"
+    ROOT_IMAGE_PATH = "%s/root.qed" % TEST_DIR
+    LEAF_IMAGE_PATH = "%s/child.qed" % TEST_DIR
+    JSON_TEST_RESULT = \
+"""{
+  "children": [
+    {
+      "children": [
+        {
+          "name": "tests/child.qed", 
+          "size": 266240
+        }
+      ], 
+      "name": "tests/root.qed", 
+      "size": 266240
+    }
+  ], 
+  "name": "/", 
+  "size": 0
+}"""
+
+    def setUp(self):
+        """ Setup the QEMU images for the test """
+        if not os.path.exists(self.TEST_DIR) and \
+           not os.path.isdir(self.TEST_DIR):
+            os.mkdir(self.TEST_DIR)    
+
+        result = subprocess.call(["qemu-img", "create", "-f", "qed",
+                                 self.ROOT_IMAGE_PATH, "1k"])
+        self.assertEqual(result, 0)
+
+        result = subprocess.call(["qemu-img", "create", "-f", "qed",
+                                 "-b", self.ROOT_IMAGE_PATH,
+                                 self.LEAF_IMAGE_PATH])
+        self.assertEqual(result, 0)
+
+    def test_get_infos_on_parent(self):
+        """ Test get_image_info on the parent snapshot """
+        (infos, result) = get_image_infos(self.ROOT_IMAGE_PATH)
+        self.assertEqual(result, True)
+        self.assertEqual(infos['image'], self.ROOT_IMAGE_PATH)
+        self.assertEqual(infos['backing_file'], None)
+        self.assertEqual(infos['file_format'], "qed")
+        self.assertEqual(infos['virtual_size'], 1024)
+        self.assertEqual(infos['cluster_size'], 65536)
+        self.assertEqual(infos['disk_size'], 266240)
+
+    def test_get_infos_on_child(self):
+        """ Test get_image_infos on the child snapshot """
+        (infos, result) = get_image_infos(self.LEAF_IMAGE_PATH)
+        self.assertEqual(result, True)
+        self.assertEqual(infos['image'], self.LEAF_IMAGE_PATH)
+        self.assertEqual(infos['backing_file'], self.ROOT_IMAGE_PATH)
+        self.assertEqual(infos['file_format'], "qed")
+        self.assertEqual(infos['virtual_size'], 1024)
+        self.assertEqual(infos['cluster_size'], 65536)
+        self.assertEqual(infos['disk_size'], 266240)
+
+    def test_json_output(self):
+        """ Test if the JSON output is good """
+        forest = Forest([self.TEST_DIR])
+        forest.build()
+        result = forest.to_json()
+        self.assertEqual(result, self.JSON_TEST_RESULT)
+
+def run_tests():
+    """ Run unit tests """
+    unittest.main()
+
+def main():
+    """ Main function """
+    sys.setrecursionlimit(1000000) # some peoples have many snapshots
+    parser = argparse.ArgumentParser(description =
+             'Build a dot(1) or D3.js JSON output representing a forest of 
qemu snapshot.')
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument('--directory', action='append',
+                       help='walk in directory and build the forest')
+    parser.add_argument('--json', action="store_true", default=False)
+    group.add_argument('--test', action='store_true', help='run the tests')
+    args = parser.parse_args()
+
+    if args.directory:
+        run_command(args.directory, args.json)
+
+    if args.test:
+        del sys.argv[1]
+        run_tests()
+
+if __name__ == '__main__':
+    main()
-- 
1.7.9.5