binman: Add command-line support for replacing entries

Add a 'replace' command to binman to permit entries to be replaced, either
individually or all at once (using a filter).

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

5 files changed, 327 insertions, 4 deletions

diff --git a/tools/binman/README b/tools/binman/README
index 5e8f9fd480..b4f6392ab7 100644
--- a/tools/binman/README
+++ b/tools/binman/README
@@ -589,9 +589,24 @@ that there is an 'fdtmap' entry in the image. For example:
     $ binman replace -i image.bin section/cbfs/u-boot
 
 which will write the contents of the file 'u-boot' from the current directory
-to the that entry. If the entry size changes, you must add the 'allow-repack'
-property to the original image before generating it (see above), otherwise you
-will get an error.
+to the that entry, compressing if necessary. If the entry size changes, you must
+add the 'allow-repack' property to the original image before generating it (see
+above), otherwise you will get an error.
+
+You can also use a particular file, in this case u-boot.bin:
+
+    $ binman replace -i image.bin section/cbfs/u-boot -f u-boot.bin
+
+It is possible to replace all files from a source directory which uses the same
+hierarchy as the entries:
+
+    $ binman replace -i image.bin -I indir
+
+Files that are missing will generate a warning.
+
+You can also replace just a selection of entries:
+
+    $ binman replace -i image.bin "*u-boot*" -I indir
 
 
 Logging
@@ -959,7 +974,6 @@ Some ideas:
 - Allow easy building of images by specifying just the board name
 - Support building an image for a board (-b) more completely, with a
   configurable build directory
-- Support updating binaries in an image (with repacking)
 - Support adding FITs to an image
 - Support for ARM Trusted Firmware (ATF)
 - Detect invalid properties in nodes
diff --git a/tools/binman/cmdline.py b/tools/binman/cmdline.py
index a43aec649e..1e38593579 100644
--- a/tools/binman/cmdline.py
+++ b/tools/binman/cmdline.py
@@ -84,6 +84,23 @@ controlled by a description in the board device tree.'''
     extract_parser.add_argument('-U', '--uncompressed', action='store_true',
         help='Output raw uncompressed data for compressed entries')
 
+    replace_parser = subparsers.add_parser('replace',
+                                           help='Replace entries in an image')
+    replace_parser.add_argument('-C', '--compressed', action='store_true',
+        help='Input data is already compressed if needed for the entry')
+    replace_parser.add_argument('-i', '--image', type=str, required=True,
+                                help='Image filename to extract')
+    replace_parser.add_argument('-f', '--filename', type=str,
+                                help='Input filename to read from')
+    replace_parser.add_argument('-F', '--fix-size', action='store_true',
+        help="Don't allow entries to be resized")
+    replace_parser.add_argument('-I', '--indir', type=str, default='',
+        help='Path to directory to use for input files')
+    replace_parser.add_argument('-m', '--map', action='store_true',
+        default=False, help='Output a map file for the updated image')
+    replace_parser.add_argument('paths', type=str, nargs='*',
+                                help='Paths within file to extract (wildcard)')
+
     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 e6722c9459..9e7587864c 100644
--- a/tools/binman/control.py
+++ b/tools/binman/control.py
@@ -253,6 +253,71 @@ def WriteEntry(image_fname, entry_path, data, do_compress=True,
 
     return image
 
+
+def ReplaceEntries(image_fname, input_fname, indir, entry_paths,
+                   do_compress=True, allow_resize=True, write_map=False):
+    """Replace the data from one or more entries from input files
+
+    Args:
+        image_fname: Image filename to process
+        input_fname: Single input ilename to use if replacing one file, None
+            otherwise
+        indir: Input directory to use (for any number of files), else None
+        entry_paths: List of entry paths to extract
+        do_compress: True if the input data is uncompressed and may need to be
+            compressed if the entry requires it, False if the data is already
+            compressed.
+        write_map: True to write a map file
+
+    Returns:
+        List of EntryInfo records that were written
+    """
+    image = Image.FromFile(image_fname)
+
+    # Replace an entry from a single file, as a special case
+    if input_fname:
+        if not entry_paths:
+            raise ValueError('Must specify an entry path to read with -f')
+        if len(entry_paths) != 1:
+            raise ValueError('Must specify exactly one entry path to write with -f')
+        entry = image.FindEntryPath(entry_paths[0])
+        data = tools.ReadFile(input_fname)
+        tout.Notice("Read %#x bytes from file '%s'" % (len(data), input_fname))
+        WriteEntryToImage(image, entry, data, do_compress=do_compress,
+                          allow_resize=allow_resize, write_map=write_map)
+        return
+
+    # Otherwise we will input from a path given by the entry path of each entry.
+    # This means that files must appear in subdirectories if they are part of
+    # a sub-section.
+    einfos = image.GetListEntries(entry_paths)[0]
+    tout.Notice("Replacing %d matching entries in image '%s'" %
+                (len(einfos), image_fname))
+
+    BeforeReplace(image, allow_resize)
+
+    for einfo in einfos:
+        entry = einfo.entry
+        if entry.GetEntries():
+            tout.Info("Skipping section entry '%s'" % entry.GetPath())
+            continue
+
+        path = entry.GetPath()[1:]
+        fname = os.path.join(indir, path)
+
+        if os.path.exists(fname):
+            tout.Notice("Write entry '%s' from file '%s'" %
+                        (entry.GetPath(), fname))
+            data = tools.ReadFile(fname)
+            ReplaceOneEntry(image, entry, data, do_compress, allow_resize)
+        else:
+            tout.Warning("Skipping entry '%s' from missing file '%s'" %
+                         (entry.GetPath(), fname))
+
+    AfterReplace(image, allow_resize=allow_resize, write_map=write_map)
+    return image
+
+
 def PrepareImagesAndDtbs(dtb_fname, select_images, update_fdt):
     """Prepare the images to be processed and select the device tree
 
@@ -420,6 +485,16 @@ def Binman(args):
             tools.FinaliseOutputDir()
         return 0
 
+    if args.cmd == 'replace':
+        try:
+            tools.PrepareOutputDir(None)
+            ReplaceEntries(args.image, args.filename, args.indir, args.paths,
+                           do_compress=not args.compressed,
+                           allow_resize=not args.fix_size, write_map=args.map)
+        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 358f095dfb..0f3b70b3bb 100644
--- a/tools/binman/ftest.py
+++ b/tools/binman/ftest.py
@@ -3047,6 +3047,195 @@ class TestFunctional(unittest.TestCase):
         data = control.ReadEntry(updated_fname, entry_name)
         self.assertEqual(expected, data)
 
+    def _SetupForReplace(self):
+        """Set up some files to use to replace entries
+
+        This generates an image, copies it to a new file, extracts all the files
+        in it and updates some of them
+
+        Returns:
+            List
+                Image filename
+                Output directory
+                Expected values for updated entries, each a string
+        """
+        data = self._DoReadFileRealDtb('143_replace_all.dts')
+
+        updated_fname = tools.GetOutputFilename('image-updated.bin')
+        tools.WriteFile(updated_fname, data)
+
+        outdir = os.path.join(self._indir, 'extract')
+        einfos = control.ExtractEntries(updated_fname, None, outdir, [])
+
+        expected1 = b'x' + U_BOOT_DATA + b'y'
+        u_boot_fname1 = os.path.join(outdir, 'u-boot')
+        tools.WriteFile(u_boot_fname1, expected1)
+
+        expected2 = b'a' + U_BOOT_DATA + b'b'
+        u_boot_fname2 = os.path.join(outdir, 'u-boot2')
+        tools.WriteFile(u_boot_fname2, expected2)
+
+        expected_text = b'not the same text'
+        text_fname = os.path.join(outdir, 'text')
+        tools.WriteFile(text_fname, expected_text)
+
+        dtb_fname = os.path.join(outdir, 'u-boot-dtb')
+        dtb = fdt.FdtScan(dtb_fname)
+        node = dtb.GetNode('/binman/text')
+        node.AddString('my-property', 'the value')
+        dtb.Sync(auto_resize=True)
+        dtb.Flush()
+
+        return updated_fname, outdir, expected1, expected2, expected_text
+
+    def _CheckReplaceMultiple(self, entry_paths):
+        """Handle replacing the contents of multiple entries
+
+        Args:
+            entry_paths: List of entry paths to replace
+
+        Returns:
+            List
+                Dict of entries in the image:
+                    key: Entry name
+                    Value: Entry object
+            Expected values for updated entries, each a string
+        """
+        updated_fname, outdir, expected1, expected2, expected_text = (
+            self._SetupForReplace())
+        control.ReplaceEntries(updated_fname, None, outdir, entry_paths)
+
+        image = Image.FromFile(updated_fname)
+        image.LoadData()
+        return image.GetEntries(), expected1, expected2, expected_text
+
+    def testReplaceAll(self):
+        """Test replacing the contents of all entries"""
+        entries, expected1, expected2, expected_text = (
+            self._CheckReplaceMultiple([]))
+        data = entries['u-boot'].data
+        self.assertEqual(expected1, data)
+
+        data = entries['u-boot2'].data
+        self.assertEqual(expected2, data)
+
+        data = entries['text'].data
+        self.assertEqual(expected_text, data)
+
+        # Check that the device tree is updated
+        data = entries['u-boot-dtb'].data
+        dtb = fdt.Fdt.FromData(data)
+        dtb.Scan()
+        node = dtb.GetNode('/binman/text')
+        self.assertEqual('the value', node.props['my-property'].value)
+
+    def testReplaceSome(self):
+        """Test replacing the contents of a few entries"""
+        entries, expected1, expected2, expected_text = (
+            self._CheckReplaceMultiple(['u-boot2', 'text']))
+
+        # This one should not change
+        data = entries['u-boot'].data
+        self.assertEqual(U_BOOT_DATA, data)
+
+        data = entries['u-boot2'].data
+        self.assertEqual(expected2, data)
+
+        data = entries['text'].data
+        self.assertEqual(expected_text, data)
+
+    def testReplaceCmd(self):
+        """Test replacing a file fron an image on the command line"""
+        self._DoReadFileRealDtb('143_replace_all.dts')
+
+        try:
+            tmpdir, updated_fname = self._SetupImageInTmpdir()
+
+            fname = os.path.join(tmpdir, 'update-u-boot.bin')
+            expected = b'x' * len(U_BOOT_DATA)
+            tools.WriteFile(fname, expected)
+
+            self._DoBinman('replace', '-i', updated_fname, 'u-boot', '-f', fname)
+            data = tools.ReadFile(updated_fname)
+            self.assertEqual(expected, data[:len(expected)])
+            map_fname = os.path.join(tmpdir, 'image-updated.map')
+            self.assertFalse(os.path.exists(map_fname))
+        finally:
+            shutil.rmtree(tmpdir)
+
+    def testReplaceCmdSome(self):
+        """Test replacing some files fron an image on the command line"""
+        updated_fname, outdir, expected1, expected2, expected_text = (
+            self._SetupForReplace())
+
+        self._DoBinman('replace', '-i', updated_fname, '-I', outdir,
+                       'u-boot2', 'text')
+
+        tools.PrepareOutputDir(None)
+        image = Image.FromFile(updated_fname)
+        image.LoadData()
+        entries = image.GetEntries()
+
+        # This one should not change
+        data = entries['u-boot'].data
+        self.assertEqual(U_BOOT_DATA, data)
+
+        data = entries['u-boot2'].data
+        self.assertEqual(expected2, data)
+
+        data = entries['text'].data
+        self.assertEqual(expected_text, data)
+
+    def testReplaceMissing(self):
+        """Test replacing entries where the file is missing"""
+        updated_fname, outdir, expected1, expected2, expected_text = (
+            self._SetupForReplace())
+
+        # Remove one of the files, to generate a warning
+        u_boot_fname1 = os.path.join(outdir, 'u-boot')
+        os.remove(u_boot_fname1)
+
+        with test_util.capture_sys_output() as (stdout, stderr):
+            control.ReplaceEntries(updated_fname, None, outdir, [])
+        self.assertIn("Skipping entry '/u-boot' from missing file",
+                      stdout.getvalue())
+
+    def testReplaceCmdMap(self):
+        """Test replacing a file fron an image on the command line"""
+        self._DoReadFileRealDtb('143_replace_all.dts')
+
+        try:
+            tmpdir, updated_fname = self._SetupImageInTmpdir()
+
+            fname = os.path.join(self._indir, 'update-u-boot.bin')
+            expected = b'x' * len(U_BOOT_DATA)
+            tools.WriteFile(fname, expected)
+
+            self._DoBinman('replace', '-i', updated_fname, 'u-boot',
+                           '-f', fname, '-m')
+            map_fname = os.path.join(tmpdir, 'image-updated.map')
+            self.assertTrue(os.path.exists(map_fname))
+        finally:
+            shutil.rmtree(tmpdir)
+
+    def testReplaceNoEntryPaths(self):
+        """Test replacing an entry without an entry path"""
+        self._DoReadFileRealDtb('143_replace_all.dts')
+        image_fname = tools.GetOutputFilename('image.bin')
+        with self.assertRaises(ValueError) as e:
+            control.ReplaceEntries(image_fname, 'fname', None, [])
+        self.assertIn('Must specify an entry path to read with -f',
+                      str(e.exception))
+
+    def testReplaceTooManyEntryPaths(self):
+        """Test extracting some entries"""
+        self._DoReadFileRealDtb('143_replace_all.dts')
+        image_fname = tools.GetOutputFilename('image.bin')
+        with self.assertRaises(ValueError) as e:
+            control.ReplaceEntries(image_fname, 'fname', None, ['a', 'b'])
+        self.assertIn('Must specify exactly one entry path to write with -f',
+                      str(e.exception))
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tools/binman/test/143_replace_all.dts b/tools/binman/test/143_replace_all.dts
new file mode 100644
index 0000000000..c5744a3c1c
--- /dev/null
+++ b/tools/binman/test/143_replace_all.dts
@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: GPL-2.0+
+
+/dts-v1/;
+
+/ {
+	#address-cells = <1>;
+	#size-cells = <1>;
+
+	binman {
+		size = <0xc00>;
+		allow-repack;
+		u-boot {
+		};
+		fdtmap {
+		};
+		u-boot2 {
+			type = "u-boot";
+		};
+		text {
+			text = "some text";
+		};
+		u-boot-dtb {
+		};
+		image-header {
+			location = "end";
+		};
+	};
+};