binman: Add an 'extract' command

It is useful to be able to extract all binaries from the image, or a
subset of them. Add a new 'extract' command to handle this.

Signed-off-by: Simon Glass <sjg@chromium.org>

4 files changed, 286 insertions, 1 deletions

diff --git a/tools/binman/README b/tools/binman/README
index 1655a9d50df..756c6a0010e 100644
--- a/tools/binman/README
+++ b/tools/binman/README
@@ -533,6 +533,30 @@ or with wildcards:
       image-header          bf8     8  image-header     bf8
 
 
+Extracting files from images
+----------------------------
+
+You can extract files from an existing firmware image created by binman,
+provided that there is an 'fdtmap' entry in the image. For example:
+
+    $ binman extract -i image.bin section/cbfs/u-boot
+
+which will write the uncompressed contents of that entry to the file 'u-boot' in
+the current directory. You can also extract to a particular file, in this case
+u-boot.bin:
+
+    $ binman extract -i image.bin section/cbfs/u-boot -f u-boot.bin
+
+It is possible to extract all files into a destination directory, which will
+put files in subdirectories matching the entry hierarchy:
+
+    $ binman extract -i image.bin -O outdir
+
+or just a selection:
+
+    $ binman extract -i image.bin "*u-boot*" -O outdir
+
+
 Logging
 -------
 
@@ -883,7 +907,6 @@ Some ideas:
 - Use of-platdata to make the information available to code that is unable
   to use device tree (such as a very small SPL image)
 - Allow easy building of images by specifying just the board name
-- Add an option to decode an image into the constituent binaries
 - Support building an image for a board (-b) more completely, with a
   configurable build directory
 - Support updating binaries in an image (with no size change / repacking)
diff --git a/tools/binman/cmdline.py b/tools/binman/cmdline.py
index 508232eabb5..a43aec649ed 100644
--- a/tools/binman/cmdline.py
+++ b/tools/binman/cmdline.py
@@ -71,6 +71,19 @@ controlled by a description in the board device tree.'''
     list_parser.add_argument('paths', type=str, nargs='*',
                              help='Paths within file to list (wildcard)')
 
+    extract_parser = subparsers.add_parser('extract',
+                                           help='Extract files from an image')
+    extract_parser.add_argument('-i', '--image', type=str, required=True,
+                                help='Image filename to extract')
+    extract_parser.add_argument('-f', '--filename', type=str,
+                                help='Output filename to write to')
+    extract_parser.add_argument('-O', '--outdir', type=str, default='',
+        help='Path to directory to use for output files')
+    extract_parser.add_argument('paths', type=str, nargs='*',
+                                help='Paths within file to extract (wildcard)')
+    extract_parser.add_argument('-U', '--uncompressed', action='store_true',
+        help='Output raw uncompressed data for compressed entries')
+
     test_parser = subparsers.add_parser('test', help='Run tests')
     test_parser.add_argument('-P', '--processes', type=int,
         help='set number of processes to use for running tests')
diff --git a/tools/binman/control.py b/tools/binman/control.py
index b244e7a0cd0..dc898be6179 100644
--- a/tools/binman/control.py
+++ b/tools/binman/control.py
@@ -118,6 +118,57 @@ def ReadEntry(image_fname, entry_path, decomp=True):
     return entry.ReadData(decomp)
 
 
+def ExtractEntries(image_fname, output_fname, outdir, entry_paths,
+                   decomp=True):
+    """Extract the data from one or more entries and write it to files
+
+    Args:
+        image_fname: Image filename to process
+        output_fname: Single output filename to use if extracting one file, None
+            otherwise
+        outdir: Output directory to use (for any number of files), else None
+        entry_paths: List of entry paths to extract
+        decomp: True to compress the entry data
+
+    Returns:
+        List of EntryInfo records that were written
+    """
+    image = Image.FromFile(image_fname)
+
+    # Output an entry to a single file, as a special case
+    if output_fname:
+        if not entry_paths:
+            raise ValueError('Must specify an entry path to write with -o')
+        if len(entry_paths) != 1:
+            raise ValueError('Must specify exactly one entry path to write with -o')
+        entry = image.FindEntryPath(entry_paths[0])
+        data = entry.ReadData(decomp)
+        tools.WriteFile(output_fname, data)
+        tout.Notice("Wrote %#x bytes to file '%s'" % (len(data), output_fname))
+        return
+
+    # Otherwise we will output to a path given by the entry path of each entry.
+    # This means that entries will appear in subdirectories if they are part of
+    # a sub-section.
+    einfos = image.GetListEntries(entry_paths)[0]
+    tout.Notice('%d entries match and will be written' % len(einfos))
+    for einfo in einfos:
+        entry = einfo.entry
+        data = entry.ReadData(decomp)
+        path = entry.GetPath()[1:]
+        fname = os.path.join(outdir, path)
+
+        # If this entry has children, create a directory for it and put its
+        # data in a file called 'root' in that directory
+        if entry.GetEntries():
+            if not os.path.exists(fname):
+                os.makedirs(fname)
+            fname = os.path.join(fname, 'root')
+        tout.Notice("Write entry '%s' to '%s'" % (entry.GetPath(), fname))
+        tools.WriteFile(fname, data)
+    return einfos
+
+
 def Binman(args):
     """The main control code for binman
 
@@ -142,6 +193,15 @@ def Binman(args):
         ListEntries(args.image, args.paths)
         return 0
 
+    if args.cmd == 'extract':
+        try:
+            tools.PrepareOutputDir(None)
+            ExtractEntries(args.image, args.filename, args.outdir, args.paths,
+                           not args.uncompressed)
+        finally:
+            tools.FinaliseOutputDir()
+        return 0
+
     # Try to figure out which device tree contains our image description
     if args.dt:
         dtb_fname = args.dt
diff --git a/tools/binman/ftest.py b/tools/binman/ftest.py
index c11dd1b85ab..709fa0adc3f 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -2446,6 +2446,43 @@ class TestFunctional(unittest.TestCase):
         data = self._RunExtractCmd('u-boot')
         self.assertEqual(U_BOOT_DATA, data)
 
+    def testExtractSection(self):
+        """Test extracting the files in a section"""
+        data = self._RunExtractCmd('section')
+        cbfs_data = data[:0x400]
+        cbfs = cbfs_util.CbfsReader(cbfs_data)
+        self.assertEqual(['u-boot', 'u-boot-dtb', ''], cbfs.files.keys())
+        dtb_data = data[0x400:]
+        dtb = self._decompress(dtb_data)
+        self.assertEqual(EXTRACT_DTB_SIZE, len(dtb))
+
+    def testExtractCompressed(self):
+        """Test extracting compressed data"""
+        data = self._RunExtractCmd('section/u-boot-dtb')
+        self.assertEqual(EXTRACT_DTB_SIZE, len(data))
+
+    def testExtractRaw(self):
+        """Test extracting compressed data without decompressing it"""
+        data = self._RunExtractCmd('section/u-boot-dtb', decomp=False)
+        dtb = self._decompress(data)
+        self.assertEqual(EXTRACT_DTB_SIZE, len(dtb))
+
+    def testExtractCbfs(self):
+        """Test extracting CBFS data"""
+        data = self._RunExtractCmd('section/cbfs/u-boot')
+        self.assertEqual(U_BOOT_DATA, data)
+
+    def testExtractCbfsCompressed(self):
+        """Test extracting CBFS compressed data"""
+        data = self._RunExtractCmd('section/cbfs/u-boot-dtb')
+        self.assertEqual(EXTRACT_DTB_SIZE, len(data))
+
+    def testExtractCbfsRaw(self):
+        """Test extracting CBFS compressed data without decompressing it"""
+        data = self._RunExtractCmd('section/cbfs/u-boot-dtb', decomp=False)
+        dtb = tools.Decompress(data, 'lzma')
+        self.assertEqual(EXTRACT_DTB_SIZE, len(dtb))
+
     def testExtractBadEntry(self):
         """Test extracting a bad section path"""
         with self.assertRaises(ValueError) as e:
@@ -2465,6 +2502,158 @@ class TestFunctional(unittest.TestCase):
         with self.assertRaises(ValueError) as e:
             control.ReadEntry(fname, 'name')
 
+    def testExtractCmd(self):
+        """Test extracting a file fron an image on the command line"""
+        self._CheckLz4()
+        self._DoReadFileRealDtb('130_list_fdtmap.dts')
+        image_fname = tools.GetOutputFilename('image.bin')
+        fname = os.path.join(self._indir, 'output.extact')
+        with test_util.capture_sys_output() as (stdout, stderr):
+            self._DoBinman('extract', '-i', image_fname, 'u-boot', '-f', fname)
+        data = tools.ReadFile(fname)
+        self.assertEqual(U_BOOT_DATA, data)
+
+    def testExtractOneEntry(self):
+        """Test extracting a single entry fron an image """
+        self._CheckLz4()
+        self._DoReadFileRealDtb('130_list_fdtmap.dts')
+        image_fname = tools.GetOutputFilename('image.bin')
+        fname = os.path.join(self._indir, 'output.extact')
+        control.ExtractEntries(image_fname, fname, None, ['u-boot'])
+        data = tools.ReadFile(fname)
+        self.assertEqual(U_BOOT_DATA, data)
+
+    def _CheckExtractOutput(self, decomp):
+        """Helper to test file output with and without decompression
+
+        Args:
+            decomp: True to decompress entry data, False to output it raw
+        """
+        def _CheckPresent(entry_path, expect_data, expect_size=None):
+            """Check and remove expected file
+
+            This checks the data/size of a file and removes the file both from
+            the outfiles set and from the output directory. Once all files are
+            processed, both the set and directory should be empty.
+
+            Args:
+                entry_path: Entry path
+                expect_data: Data to expect in file, or None to skip check
+                expect_size: Size of data to expect in file, or None to skip
+            """
+            path = os.path.join(outdir, entry_path)
+            data = tools.ReadFile(path)
+            os.remove(path)
+            if expect_data:
+                self.assertEqual(expect_data, data)
+            elif expect_size:
+                self.assertEqual(expect_size, len(data))
+            outfiles.remove(path)
+
+        def _CheckDirPresent(name):
+            """Remove expected directory
+
+            This gives an error if the directory does not exist as expected
+
+            Args:
+                name: Name of directory to remove
+            """
+            path = os.path.join(outdir, name)
+            os.rmdir(path)
+
+        self._DoReadFileRealDtb('130_list_fdtmap.dts')
+        image_fname = tools.GetOutputFilename('image.bin')
+        outdir = os.path.join(self._indir, 'extract')
+        einfos = control.ExtractEntries(image_fname, None, outdir, [], decomp)
+
+        # Create a set of all file that were output (should be 9)
+        outfiles = set()
+        for root, dirs, files in os.walk(outdir):
+            outfiles |= set([os.path.join(root, fname) for fname in files])
+        self.assertEqual(9, len(outfiles))
+        self.assertEqual(9, len(einfos))
+
+        image = control.images['image']
+        entries = image.GetEntries()
+
+        # Check the 9 files in various ways
+        section = entries['section']
+        section_entries = section.GetEntries()
+        cbfs_entries = section_entries['cbfs'].GetEntries()
+        _CheckPresent('u-boot', U_BOOT_DATA)
+        _CheckPresent('section/cbfs/u-boot', U_BOOT_DATA)
+        dtb_len = EXTRACT_DTB_SIZE
+        if not decomp:
+            dtb_len = cbfs_entries['u-boot-dtb'].size
+        _CheckPresent('section/cbfs/u-boot-dtb', None, dtb_len)
+        if not decomp:
+            dtb_len = section_entries['u-boot-dtb'].size
+        _CheckPresent('section/u-boot-dtb', None, dtb_len)
+
+        fdtmap = entries['fdtmap']
+        _CheckPresent('fdtmap', fdtmap.data)
+        hdr = entries['image-header']
+        _CheckPresent('image-header', hdr.data)
+
+        _CheckPresent('section/root', section.data)
+        cbfs = section_entries['cbfs']
+        _CheckPresent('section/cbfs/root', cbfs.data)
+        data = tools.ReadFile(image_fname)
+        _CheckPresent('root', data)
+
+        # There should be no files left. Remove all the directories to check.
+        # If there are any files/dirs remaining, one of these checks will fail.
+        self.assertEqual(0, len(outfiles))
+        _CheckDirPresent('section/cbfs')
+        _CheckDirPresent('section')
+        _CheckDirPresent('')
+        self.assertFalse(os.path.exists(outdir))
+
+    def testExtractAllEntries(self):
+        """Test extracting all entries"""
+        self._CheckLz4()
+        self._CheckExtractOutput(decomp=True)
+
+    def testExtractAllEntriesRaw(self):
+        """Test extracting all entries without decompressing them"""
+        self._CheckLz4()
+        self._CheckExtractOutput(decomp=False)
+
+    def testExtractSelectedEntries(self):
+        """Test extracting some entries"""
+        self._CheckLz4()
+        self._DoReadFileRealDtb('130_list_fdtmap.dts')
+        image_fname = tools.GetOutputFilename('image.bin')
+        outdir = os.path.join(self._indir, 'extract')
+        einfos = control.ExtractEntries(image_fname, None, outdir,
+                                        ['*cb*', '*head*'])
+
+        # File output is tested by testExtractAllEntries(), so just check that
+        # the expected entries are selected
+        names = [einfo.name for einfo in einfos]
+        self.assertEqual(names,
+                         ['cbfs', 'u-boot', 'u-boot-dtb', 'image-header'])
+
+    def testExtractNoEntryPaths(self):
+        """Test extracting some entries"""
+        self._CheckLz4()
+        self._DoReadFileRealDtb('130_list_fdtmap.dts')
+        image_fname = tools.GetOutputFilename('image.bin')
+        with self.assertRaises(ValueError) as e:
+            control.ExtractEntries(image_fname, 'fname', None, [])
+        self.assertIn('Must specify an entry path to write with -o',
+                      str(e.exception))
+
+    def testExtractTooManyEntryPaths(self):
+        """Test extracting some entries"""
+        self._CheckLz4()
+        self._DoReadFileRealDtb('130_list_fdtmap.dts')
+        image_fname = tools.GetOutputFilename('image.bin')
+        with self.assertRaises(ValueError) as e:
+            control.ExtractEntries(image_fname, 'fname', None, ['a', 'b'])
+        self.assertIn('Must specify exactly one entry path to write with -o',
+                      str(e.exception))
+
 
 if __name__ == "__main__":
     unittest.main()