1 files changed, 99 insertions, 0 deletions
diff --git a/images/uhdimgs.py b/images/uhdimgs.py
new file mode 100644
index 000000000..77da36c0d
--- /dev/null
+++ b/images/uhdimgs.py
@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+#
+# Copyright 2014 Ettus Research LLC
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+"""
+Utility module for packaging and handling UHD binary images.
+"""
+
+import re
+import os
+import sys
+import hashlib
+
+_DEFAULT_BASE_URL = "http://files.ettus.com/binaries/images"
+_DEFAULT_BUFFER_SIZE      = 8192
+_CMAKE_MAIN_FILE = "../host/CMakeLists.txt"
+_CMAKE_IMG_TPL = """
+# This section is written automatically by /images/create_imgs_package.py
+# Any manual changes in here will be overwritten.
+SET(UHD_IMAGES_MD5SUM "{}")
+SET(UHD_IMAGES_DOWNLOAD_SRC "{}")
+"""
+
+def get_cmake_main_file():
+    """ Returns a path to the CMakeLists.txt file that contains the image info. """
+    return _CMAKE_MAIN_FILE
+
+def get_base_url():
+    """ Returns the base URL """
+    if os.environ.get("UHD_IMAGES_BASE_URL") is not None and os.environ.get("UHD_IMAGES_BASE_URL") != "":
+        return os.environ.get("UHD_IMAGES_BASE_URL")
+    else:
+        return _DEFAULT_BASE_URL
+
+def base_url_is_local(base_url):
+    """ Returns true if the base URL is actually a http URL
+    instead of a local path. """
+    return not (base_url.find('http') == 0)
+
+def get_images_dir():
+    """
+    Returns the absolute position of the images/ subdir
+    in the UHD source tree.
+    """
+    return os.path.dirname(__file__)
+
+def update_main_cmake_file(md5, zipfilename):
+    """
+    Update the section in host/CMakeLists.txt that contains
+    the ZIP filename and the MD5 hash.
+    """
+    cmakef = open(_CMAKE_MAIN_FILE, 'r').read()
+    new_section = _CMAKE_IMG_TPL.format(md5, zipfilename)
+    regex = re.compile("(?<={{{IMG_SECTION)(.*)(?=#}}})", flags=re.MULTILINE|re.DOTALL)
+    cmakef = regex.sub(new_section, cmakef, count=1)
+    open(_CMAKE_MAIN_FILE, 'w').write(cmakef)
+
+def get_total_md5(img_dir):
+    """ Creates an md5sum of everything in the images directory """
+    def _update_md5_for_dir_recursive(dir_root, md5_obj):
+        for (root, dirnames, filenames) in os.walk(dir_root):
+            for filename in filenames:
+                md5_obj.update(open(os.path.join(root, filename), 'rb').read())
+                sys.stdout.write('.')
+                sys.stdout.flus()
+            for dirname in dirnames:
+                _update_md5_for_dir_recursive(os.path.join(root, dirname), md5_obj)
+    md5 = hashlib.md5()
+    _update_md5_for_dir_recursive(img_dir, md5)
+    print ""
+    return md5.hexdigest()
+
+def md5_checksum(filePath):
+    """ Return MD5 checksum of a single file. """
+    try:
+        with open(filePath, 'rb') as fh:
+            m = hashlib.md5()
+            while True:
+                data = fh.read(_DEFAULT_BUFFER_SIZE)
+                if not data:
+                    break
+                m.update(data)
+            return m.hexdigest()
+    except Exception, e:
+        print "Failed to calculated MD5 sum of: %s (%s)" % (filePath, e)
+        raise e