Implement support for --chmod and --chown rsync flags (supersedes PR #46)

afl_utils/afl_sync.py

@@ -24,20 +24,20 @@
 
 _rsync_default_options = ['-racz']
 
-
 class AflBaseSync(object):
-    def __init__(self, server_config, fuzzer_config):
+    def __init__(self, server_config, fuzzer_config, rsync_config):
         self.server_config = server_config
         self.fuzzer_config = fuzzer_config
+        self.rsync_config  = rsync_config
 
 
 class AflRsync(AflBaseSync):
-    def __init__(self, server_config, fuzzer_config):
+    def __init__(self, server_config, fuzzer_config, rsync_config):
         # default excludes
         self.__excludes = ['*.cur_input']
-        super(AflRsync, self).__init__(server_config, fuzzer_config)
+        super(AflRsync, self).__init__(server_config, fuzzer_config, rsync_config)
 
-    def __prepare_rsync_commandline(self, local_path, remote_path, rsync_options=list(_rsync_default_options),
+    def __prepare_rsync_commandline(self, local_path, remote_path, rsync_options,
                                     rsync_excludes=list([]), rsync_get=False):
         cmd = ['rsync']
 
@@ -66,11 +66,11 @@ def __get_fuzzers(self):
         fuzzers = (fuzzer for fuzzer in fuzzers if not fuzzer.endswith('.sync'))
         return fuzzers
 
-    def rsync_put(self, local_path, remote_path, rsync_options=list(_rsync_default_options), rsync_excludes=list([])):
+    def rsync_put(self, local_path, remote_path, rsync_options, rsync_excludes=list([])):
         cmd = self.__prepare_rsync_commandline(local_path, remote_path, rsync_options, rsync_excludes)
         return self.__invoke_rsync(cmd)
 
-    def rsync_get(self, remote_path, local_path, rsync_options=list(_rsync_default_options), rsync_excludes=list([])):
+    def rsync_get(self, remote_path, local_path, rsync_options, rsync_excludes=list([])):
         cmd = self.__prepare_rsync_commandline(local_path, remote_path, rsync_options, rsync_excludes, True)
         return self.__invoke_rsync(cmd)
 
@@ -93,15 +93,15 @@ def push(self):
             local_path = os.path.join(self.fuzzer_config['sync_dir'], f)
             remote_path = os.path.join(self.server_config['remote_path'], f)
             print_ok('Pushing {} -> {}.sync'.format(local_path, remote_path))
-            self.rsync_put(local_path, remote_path, rsync_excludes=excludes)
+            self.rsync_put(local_path, remote_path, self.rsync_config['put'], rsync_excludes=excludes)
 
     def pull(self):
         fuzzers = self.__get_fuzzers()
 
         local_path = self.fuzzer_config['sync_dir']
         remote_path = self.server_config['remote_path']
 
-        options = list(_rsync_default_options)
+        options = self.rsync_config['get'][:]
         excludes = self.__excludes
 
         # exclude our previously pushed fuzzer states from being pulled again
@@ -150,6 +150,12 @@ def main(argv):
                         help='Source afl synchronisation directory containing state directories of afl instances.')
     parser.add_argument('dst_storage_dir',
                         help='Destination directory used as fuzzer state storage. This shouldn\'t be an afl sync dir!')
+    parser.add_argument('--chmod',
+                        help='Affect destination\'s file and directory permissions, e.g. --chmod=g+rw to add '
+                             'read/write group permissions.', metavar="PERMS")
+    parser.add_argument('--chown',
+                        help='Affect destination\'s file and directory user and group, e.g. --chown=foo:bar to '
+                        'let the files be owned by user foo and group bar.', metavar="USER:GROUP")
     parser.add_argument('-S', '--session', dest='session', default=None,
                         help='Name of an afl-multicore session. If provided, only fuzzers belonging to '
                              'the specified session will be synced with the destination. Otherwise state '
@@ -163,6 +169,21 @@ def main(argv):
         print_err('Sorry, unknown command requested!')
         sys.exit(1)
 
+    rsync_put_options = _rsync_default_options[:]
+    rsync_get_options = _rsync_default_options[:]
+
+    if args.chmod or args.chown:
+        # these arguments are meaningless with pull since they should only
+        # affect the remote side
+        if args.cmd == 'pull':
+            print_warn('--chmod and --chown have no effect with pull and will be ignored.')
+
+        if args.chmod:
+            rsync_put_options.append("--chmod={}".format(args.chmod))
+
+        if args.chown:
+            rsync_put_options.append(["--protect-args", "--chown={}".format(args.chown)])
+
     if not os.path.exists(args.src_sync_dir):
         if args.cmd in ['pull', 'sync']:
             print_warn('Local afl sync dir does not exist! Will create it for you!')
@@ -182,7 +203,12 @@ def main(argv):
         'exclude_hangs':    False,
     }
 
-    rsyncEngine = AflRsync(server_config, fuzzer_config)
+    rsync_config = {
+        'get': rsync_get_options,
+        'put': rsync_put_options,
+    }
+
+    rsyncEngine = AflRsync(server_config, fuzzer_config, rsync_config)
 
     if args.cmd == 'push':
         rsyncEngine.push()

tests/test_afl_sync.py

@@ -77,13 +77,19 @@ def test_afl_rsync_init(self):
             'exclude_hangs': True,
         }
 
-        afl_rsync = AflRsync(server_config, fuzzer_config)
+        rsync_config = {
+            'get': afl_sync._rsync_default_options[:],
+            'put': afl_sync._rsync_default_options[:],
+        }
+
+        afl_rsync = AflRsync(server_config, fuzzer_config, rsync_config)
 
         self.assertDictEqual(server_config, afl_rsync.server_config)
         self.assertDictEqual(fuzzer_config, afl_rsync.fuzzer_config)
+        self.assertDictEqual(rsync_config, afl_rsync.rsync_config)
 
     def test_afl_rsync_prepare_sync_command(self):
-        afl_rsync = AflRsync(None, None)
+        afl_rsync = AflRsync(None, None, None)
 
         expected_put_cmdline = [
             'rsync',
@@ -102,15 +108,17 @@ def test_afl_rsync_prepare_sync_command(self):
         ]
 
         self.assertListEqual(expected_put_cmdline, afl_rsync._AflRsync__prepare_rsync_commandline('src', 'dst',
+                                                                                                  list(afl_sync._rsync_default_options),
                                                                                                   rsync_excludes=[
                                                                                                       'exclude']))
         self.assertListEqual(expected_get_cmdline, afl_rsync._AflRsync__prepare_rsync_commandline('src', 'dst',
+                                                                                    list(afl_sync._rsync_default_options),
                                                                                     rsync_excludes=['exclude'],
                                                                                     rsync_get=True))
 
     def test_afl_rsync_invoke_rsync(self):
         rsync_cmdline = ['rsync', '--help']
-        afl_rsync = AflRsync(None, None)
+        afl_rsync = AflRsync(None, None, None)
 
         self.assertTrue(afl_rsync._AflRsync__invoke_rsync(rsync_cmdline))
         self.assertFalse(afl_rsync._AflRsync__invoke_rsync(['rsync']))
@@ -130,16 +138,21 @@ def test_afl_rsync_get_fuzzers(self):
             'invalid_fuzz001'
         ]
 
-        afl_rsync = AflRsync(None, fuzzer_config)
+        afl_rsync = AflRsync(None, fuzzer_config, None)
         self.assertListEqual(sorted(expected_fuzzers), sorted(afl_rsync._AflRsync__get_fuzzers()))
 
     def test_afl_rsync_put(self):
         local_path = 'testdata/sync/fuzz000'
         remote_path = 'testdata/rsync_tmp_store/fuzz000'
         excludes = ['crashes*/', 'hangs*/']
 
-        afl_rsync = AflRsync(None, None)
-        self.assertTrue(afl_rsync.rsync_put(local_path, remote_path, rsync_excludes=excludes))
+        rsync_config = {
+            'get': afl_sync._rsync_default_options[:],
+            'put': afl_sync._rsync_default_options[:],
+        }
+
+        afl_rsync = AflRsync(None, None, rsync_config)
+        self.assertTrue(afl_rsync.rsync_put(local_path, remote_path, afl_rsync.rsync_config['put'], rsync_excludes=excludes))
         self.assertTrue(os.path.exists(remote_path + '.sync/fuzzer_stats'))
         self.assertTrue(os.path.exists(remote_path + '.sync/.cur_input'))
         self.assertFalse(os.path.exists(remote_path + '.sync/crashes'))
@@ -150,8 +163,13 @@ def test_afl_rsync_get(self):
         remote_path = 'testdata/sync/fuzz000'
         excludes = ['crashes*/', 'hangs*/']
 
-        afl_rsync = AflRsync(None, None)
-        self.assertTrue(afl_rsync.rsync_get(remote_path, local_path, rsync_excludes=excludes))
+        rsync_config = {
+            'get': afl_sync._rsync_default_options[:],
+            'put': afl_sync._rsync_default_options[:],
+        }
+
+        afl_rsync = AflRsync(None, None, rsync_config)
+        self.assertTrue(afl_rsync.rsync_get(remote_path, local_path, afl_rsync.rsync_config['get'], rsync_excludes=excludes))
         self.assertTrue(os.path.exists(local_path + '/fuzzer_stats'))
         self.assertFalse(os.path.exists(local_path + '/crashes'))
         self.assertFalse(os.path.exists(local_path + '/hangs'))
@@ -168,12 +186,16 @@ def test_afl_rsync_push(self):
             'exclude_hangs': True,
         }
 
-        afl_rsync = AflRsync(server_config, fuzzer_config)
+        rsync_config = {
+            'get': afl_sync._rsync_default_options[:],
+            'put': afl_sync._rsync_default_options[:],
+        }
+
+        afl_rsync = AflRsync(server_config, fuzzer_config, rsync_config)
         self.assertIsNone(afl_rsync.push())
         self.assertTrue(os.path.exists('testdata/rsync_output_push/fuzz000.sync'))
         self.assertFalse(os.path.exists('testdata/rsync_output_push/fuzz000.sync/.cur_input'))
         self.assertTrue(os.path.exists('testdata/rsync_output_push/fuzz001.sync'))
-        self.assertFalse(os.path.exists('testdata/rsync_output_push/fuzz000.sync/.cur_input'))
         self.assertFalse(os.path.exists('testdata/rsync_output_push/fuzz002.sync'))
         self.assertFalse(os.path.exists('testdata/rsync_output_push/fuzz002.sync.sync'))
         self.assertFalse(os.path.exists('testdata/rsync_output_push/invalid_fuzz000.sync'))
@@ -191,7 +213,12 @@ def test_afl_rsync_pull_session(self):
             'exclude_hangs': True,
         }
 
-        afl_rsync = AflRsync(server_config, fuzzer_config)
+        rsync_config = {
+            'get': afl_sync._rsync_default_options[:],
+            'put': afl_sync._rsync_default_options[:],
+        }
+
+        afl_rsync = AflRsync(server_config, fuzzer_config, rsync_config)
         self.assertIsNone(afl_rsync.pull())
         self.assertTrue(os.path.exists('testdata/sync/other_fuzz000.sync'))
         self.assertTrue(os.path.exists('testdata/sync/other_fuzz000.sync/crashes'))
@@ -214,7 +241,12 @@ def test_afl_rsync_pull_all(self):
             'exclude_hangs': True,
         }
 
-        afl_rsync = AflRsync(server_config, fuzzer_config)
+        rsync_config = {
+            'get': afl_sync._rsync_default_options[:],
+            'put': afl_sync._rsync_default_options[:],
+        }
+
+        afl_rsync = AflRsync(server_config, fuzzer_config, rsync_config)
         self.assertIsNone(afl_rsync.pull())
         self.assertTrue(os.path.exists('testdata/sync/other_fuzz000.sync'))
         self.assertTrue(os.path.exists('testdata/sync/other_fuzz001.sync'))
@@ -236,7 +268,12 @@ def test_afl_rsync_sync(self):
             'exclude_hangs': True,
         }
 
-        afl_rsync = AflRsync(server_config, fuzzer_config)
+        rsync_config = {
+            'get': afl_sync._rsync_default_options[:],
+            'put': afl_sync._rsync_default_options[:],
+        }
+
+        afl_rsync = AflRsync(server_config, fuzzer_config, rsync_config)
         self.assertIsNone(afl_rsync.sync())
 
         # pull assertions