From 14b8721a55c38082d1f2d62d7c4c1ff08acb1217 Mon Sep 17 00:00:00 2001 From: Brian Wylie Date: Sat, 2 Aug 2014 17:08:36 -0600 Subject: [PATCH] improve temp_dir context manager --- workbench/workers/pcap_bro.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/workbench/workers/pcap_bro.py b/workbench/workers/pcap_bro.py index b9612ed..7e2987f 100644 --- a/workbench/workers/pcap_bro.py +++ b/workbench/workers/pcap_bro.py @@ -63,8 +63,7 @@ def execute(self, input_data): script_path = self.get_bro_script_path() # Create a temporary directory - with self.make_temp_directory() as temp_dir: - os.chdir(temp_dir) + with self.goto_temp_directory() as temp_dir: # Get the pcap inputs (filenames) print 'pcap_bro: Setting up PCAP inputs...' @@ -131,16 +130,17 @@ def subprocess_manager(self, exec_args): raise RuntimeError('%s\npcap_bro had returncode: %d' % (exec_args, sp.returncode)) @contextlib.contextmanager - def make_temp_directory(self): - ''' Bro temporary directory context manager ''' + def goto_temp_directory(self): + previousDir = os.getcwd() temp_dir = tempfile.mkdtemp() + os.chdir(temp_dir) try: yield temp_dir finally: + # Change back to original directory + os.chdir(previousDir) # Remove the directory/files shutil.rmtree(temp_dir) - # Change back to original directory - os.chdir(self.orig_dir) def __del__(self): ''' Class Cleanup '''