diff --git a/afl_utils/afl_sync.py b/afl_utils/afl_sync.py index 8057606..94f1433 100644 --- a/afl_utils/afl_sync.py +++ b/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,7 +93,7 @@ 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() @@ -101,7 +101,7 @@ def pull(self): 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() diff --git a/tests/test_afl_sync.py b/tests/test_afl_sync.py index 746aa1b..b7ae2f5 100644 --- a/tests/test_afl_sync.py +++ b/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,7 +138,7 @@ 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): @@ -138,8 +146,13 @@ def test_afl_rsync_put(self): 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