cli: Support writing timestamp to Manifest

2 files changed, 34 insertions, 6 deletions

diff --git a/gemato/cli.py b/gemato/cli.py
index a73d875..970b4a8 100644
--- a/gemato/cli.py
+++ b/gemato/cli.py
@@ -133,12 +133,23 @@ def do_update(args, argp):
             relpath = os.path.relpath(p, os.path.dirname(tlm))
             if relpath == '.':
                 relpath = ''
+            if args.timestamp and relpath != '':
+                argp.error('Timestamp can only be updated if doing full-tree update')
+
             try:
+                start_ts = datetime.datetime.utcnow()
                 m.update_entries_for_directory(relpath)
 
-                ts = m.find_timestamp()
-                if ts is not None:
-                    ts.ts = datetime.datetime.utcnow()
+                # write TIMESTAMP if requested, or if already there
+                if relpath != '':
+                    # skip timestamp if not doing full update
+                    pass
+                elif args.timestamp:
+                    m.set_timestamp(start_ts)
+                else:
+                    ts = m.find_timestamp()
+                    if ts is not None:
+                        ts.ts = start_ts
 
                 m.save_manifests(**save_kwargs)
             except gemato.exceptions.ManifestCrossDevice as e:
@@ -197,11 +208,12 @@ def do_create(args, argp):
                 argp.error('--hashes must be specified if not implied by --profile')
 
             try:
+                start_ts = datetime.datetime.utcnow()
                 m.update_entries_for_directory()
 
-                ts = m.find_timestamp()
-                if ts is not None:
-                    ts.ts = datetime.datetime.utcnow()
+                # write TIMESTAMP if requested, or if already there
+                if args.timestamp:
+                    m.set_timestamp(start_ts)
 
                 m.save_manifests(**save_kwargs)
             except gemato.exceptions.ManifestCrossDevice as e:
@@ -265,6 +277,8 @@ def main(argv):
     signgroup.add_argument('-S', '--no-sign', action='store_false',
             dest='sign',
             help='Disable signing the top-level Manifest')
+    update.add_argument('-t', '--timestamp', action='store_true',
+            help='Include TIMESTAMP entry in Manifest')
     update.set_defaults(func=do_update)
 
     create = subp.add_parser('create',
@@ -292,6 +306,8 @@ def main(argv):
     signgroup.add_argument('-S', '--no-sign', action='store_false',
             dest='sign',
             help='Disable signing the top-level Manifest')
+    create.add_argument('-t', '--timestamp', action='store_true',
+            help='Include TIMESTAMP entry in Manifest')
     create.set_defaults(func=do_create)
 
     vals = argp.parse_args(argv[1:])
diff --git a/tests/test_recursiveloader.py b/tests/test_recursiveloader.py
index 2835b7f..3d51d8b 100644
--- a/tests/test_recursiveloader.py
+++ b/tests/test_recursiveloader.py
@@ -959,6 +959,18 @@ DATA test 0 MD5 d41d8cd98f00b204e9800998ecf8427e
         self.assertEqual(m.find_timestamp().ts,
                 datetime.datetime(2010, 7, 7, 7, 7, 7))
 
+    def test_cli_update_with_timestamp(self):
+        self.assertEqual(
+                gemato.cli.main(['gemato', 'update',
+                    '--hashes=SHA256 SHA512',
+                    '--timestamp',
+                    self.dir]),
+                0)
+
+        m = gemato.recursiveloader.ManifestRecursiveLoader(
+            os.path.join(self.dir, 'Manifest'))
+        self.assertIsNotNone(m.find_timestamp())
+
 
 class DuplicateManifestFileEntryTest(TempDirTestCase):
     """