import os from pathlib import Path import sys import tempfile import unittest from unittest import mock ROOT = Path(__file__).resolve().parents[1] APP_DIR = ROOT / 'app' sys.path.insert(0, str(APP_DIR)) import janitor from process_identity import current_process_identity from runtime_security import atomic_write_private_json, ensure_private_directory, harden_private_file class JanitorBoundedTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.root = os.path.join(self.temp.name, 'work') ensure_private_directory(self.root, reject_reparse=True) self.identity = current_process_identity().as_dict() def tearDown(self): self.temp.cleanup() def candidate(self, name='trufflehog-run-fixture'): path = os.path.join(self.root, name) ensure_private_directory(path, reject_reparse=True) atomic_write_private_json(os.path.join(path, janitor.MARKER_NAME), { 'schema': 2, 'owner_pid': self.identity['pid'], 'owner_creation_time': self.identity['creation_time'], 'owner_executable': self.identity['executable'], 'parent_pid': self.identity['pid'], 'parent_creation_time': self.identity['creation_time'], 'parent_executable': self.identity['executable'], 'created_at': '2020-01-01T00:00:00+00:00', 'root_kind': 'work', 'relative_path': name, }) return path def test_unknown_identity_retains_exact_tree(self): path = self.candidate() with mock.patch.object(janitor, 'exact_process_identity_state', return_value='unknown'): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, ) self.assertEqual(report['removed'], 0) self.assertTrue(os.path.isdir(path)) def test_dead_child_owner_allows_cleanup_while_distinct_parent_lives(self): path = self.candidate() marker_path = os.path.join(path, janitor.MARKER_NAME) marker = janitor.read_private_json(marker_path) marker['owner_pid'] = int(marker['parent_pid']) + 1000 marker['owner_creation_time'] = 'distinct-child-creation' atomic_write_private_json(marker_path, marker) child = os.path.join(path, 'result.txt') Path(child).write_text('fixture', encoding='ascii') harden_private_file(child) with mock.patch.object( janitor, 'exact_process_identity_state', side_effect=('dead', 'alive'), ): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, ) self.assertEqual(report['removed'], 1) self.assertFalse(os.path.exists(path)) def test_live_parent_retains_parent_owned_directory(self): path = self.candidate() with mock.patch.object( janitor, 'exact_process_identity_state', side_effect=('dead', 'alive'), ): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, ) self.assertEqual(report['removed'], 0) self.assertTrue(os.path.isdir(path)) def test_entry_budget_stops_immediately_and_leaves_marker(self): path = self.candidate() for index in range(4): child = os.path.join(path, f'file-{index}.txt') Path(child).write_text('fixture', encoding='ascii') harden_private_file(child) budget = janitor.JanitorBudget( max_candidates=1, max_entries=1, max_bytes=1024, max_seconds=30, max_depth=64, ) with mock.patch.object(janitor, 'exact_process_identity_state', return_value='dead'): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, budget=budget, ) self.assertTrue(report['exhausted']) self.assertTrue(os.path.isfile(os.path.join(path, janitor.MARKER_NAME))) def test_single_oversized_file_makes_forward_progress(self): path = self.candidate() child = os.path.join(path, 'oversized.pack') Path(child).write_bytes(b'x' * 2048) harden_private_file(child) budget = janitor.JanitorBudget( max_candidates=1, max_entries=10, max_bytes=1024, max_seconds=30, max_depth=64, ) with mock.patch.object(janitor, 'exact_process_identity_state', return_value='dead'): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, budget=budget, ) self.assertTrue(report['exhausted']) self.assertEqual(report['removed'], 1) self.assertFalse(os.path.exists(path)) def test_nested_oversized_file_makes_forward_progress(self): path = self.candidate() nested = os.path.join(path, 'objects', 'pack') ensure_private_directory(nested, reject_reparse=True) child = os.path.join(nested, 'oversized.pack') Path(child).write_bytes(b'x' * 2048) harden_private_file(child) budget = janitor.JanitorBudget( max_candidates=1, max_entries=10, max_bytes=1024, max_seconds=30, max_depth=64, ) with mock.patch.object(janitor, 'exact_process_identity_state', return_value='dead'): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, budget=budget, ) self.assertTrue(report['exhausted']) self.assertEqual(report['removed'], 1) self.assertFalse(os.path.exists(path)) def test_depth_budget_retains_deep_tree(self): path = self.candidate() nested = os.path.join(path, 'one') ensure_private_directory(nested, reject_reparse=True) nested = os.path.join(nested, 'two') ensure_private_directory(nested, reject_reparse=True) budget = janitor.JanitorBudget(max_depth=1) with mock.patch.object(janitor, 'exact_process_identity_state', return_value='dead'): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, budget=budget, ) self.assertTrue(report['exhausted']) self.assertTrue(os.path.isfile(os.path.join(path, janitor.MARKER_NAME))) def test_foreign_executable_and_markerless_paths_are_retained(self): path = self.candidate() markerless = os.path.join(self.root, 'trufflehog-run-markerless') ensure_private_directory(markerless, reject_reparse=True) with mock.patch.object(janitor, 'exact_process_identity_state', return_value='dead'): report = janitor.run_janitor_pass( self.root, [os.path.join(self.temp.name, 'foreign.exe')], minimum_age_sec=0, ) self.assertEqual(report['removed'], 0) self.assertTrue(os.path.isdir(path)) self.assertTrue(os.path.isdir(markerless)) def test_dead_exact_identities_delete_marker_last(self): path = self.candidate() child = os.path.join(path, 'result.txt') Path(child).write_text('fixture', encoding='ascii') harden_private_file(child) with mock.patch.object(janitor, 'exact_process_identity_state', return_value='dead'): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, ) self.assertEqual(report['removed'], 1) self.assertFalse(os.path.exists(path)) def test_huge_first_layout_cannot_starve_later_eligible_object(self): first = os.path.join(self.root, 'first') later = os.path.join(self.root, 'later') ensure_private_directory(first, reject_reparse=True) ensure_private_directory(later, reject_reparse=True) for index in range(2000): os.mkdir(os.path.join(first, f'tmp-retained-{index:05d}')) eligible = os.path.join(later, 'tmp-eligible') ensure_private_directory(eligible, reject_reparse=True) atomic_write_private_json(os.path.join(eligible, janitor.MARKER_NAME), { 'schema': 2, 'owner_pid': self.identity['pid'], 'owner_creation_time': self.identity['creation_time'], 'owner_executable': self.identity['executable'], 'parent_pid': self.identity['pid'], 'parent_creation_time': self.identity['creation_time'], 'parent_executable': self.identity['executable'], 'created_at': '2020-01-01T00:00:00+00:00', 'root_kind': 'work', 'relative_path': 'later/tmp-eligible', }) cursor_path = os.path.join(self.temp.name, 'state', 'janitor.cursor.json') cursor = janitor.JanitorCursorStore(cursor_path, self.root) layouts = ( ('work', 'first', ('tmp-',)), ('work', 'later', ('tmp-',)), ) try: with mock.patch.object(janitor, 'APPROVED_LAYOUTS', layouts), \ mock.patch.object(janitor, 'exact_process_identity_state', return_value='dead'): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, budget=janitor.JanitorBudget( max_candidates=4, max_enumerated=20, max_seconds=30, ), cursor_store=cursor, ) finally: cursor.close() self.assertGreaterEqual(report['enumerated'], 2) self.assertLessEqual(report['enumerated'], 20) self.assertEqual(report['removed'], 1) self.assertFalse(os.path.exists(eligible)) self.assertTrue(os.path.isfile(cursor_path)) def test_wrapped_small_layout_cannot_starve_seeking_large_layout(self): large = os.path.join(self.root, 'large') small = os.path.join(self.root, 'small') ensure_private_directory(large, reject_reparse=True) ensure_private_directory(small, reject_reparse=True) for index in range(20): ensure_private_directory(os.path.join(large, f'skip-{index:02d}'), reject_reparse=True) ensure_private_directory(os.path.join(large, 'tmp-anchor'), reject_reparse=True) eligible = os.path.join(large, 'tmp-eligible') ensure_private_directory(eligible, reject_reparse=True) atomic_write_private_json(os.path.join(eligible, janitor.MARKER_NAME), { 'schema': 2, 'owner_pid': self.identity['pid'], 'owner_creation_time': self.identity['creation_time'], 'owner_executable': self.identity['executable'], 'parent_pid': self.identity['pid'], 'parent_creation_time': self.identity['creation_time'], 'parent_executable': self.identity['executable'], 'created_at': '2020-01-01T00:00:00+00:00', 'root_kind': 'work', 'relative_path': 'large/tmp-eligible', }) ensure_private_directory(os.path.join(small, 'tmp-retained'), reject_reparse=True) layouts = ( ('work', 'large', ('tmp-',)), ('work', 'small', ('tmp-',)), ) cursor = janitor._MemoryCursorStore() large_name = janitor._layout_cursor_name(self.root, 'large', ('tmp-',)) cursor.state['layouts'][large_name] = {'last_name': 'tmp-anchor', 'wrap_count': 0} with mock.patch.object(janitor, 'APPROVED_LAYOUTS', layouts), \ mock.patch.object(janitor, 'exact_process_identity_state', return_value='dead'): report = janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, budget=janitor.JanitorBudget( max_candidates=4, max_enumerated=100, max_seconds=30, ), cursor_store=cursor, ) self.assertEqual(report['removed'], 1) self.assertFalse(os.path.exists(eligible)) def test_persistent_cursor_flushes_once_per_pass(self): self.candidate('trufflehog-run-one') self.candidate('trufflehog-run-two') cursor_path = os.path.join(self.temp.name, 'state', 'janitor.cursor.json') cursor = janitor.JanitorCursorStore(cursor_path, self.root) try: with mock.patch.object(cursor, '_save', wraps=cursor._save) as save, \ mock.patch.object(janitor, 'exact_process_identity_state', return_value='unknown'): janitor.run_janitor_pass( self.root, [self.identity['executable']], minimum_age_sec=0, cursor_store=cursor, ) self.assertEqual(save.call_count, 1) finally: cursor.close() if __name__ == '__main__': unittest.main()