diff --git a/mercury/app.py b/mercury/app.py index 6e613d92..2d44560e 100644 --- a/mercury/app.py +++ b/mercury/app.py @@ -2,7 +2,6 @@ import json import os import shutil -import subprocess import sys import tempfile @@ -27,16 +26,12 @@ def send_rabbit_msg(msg, channel, exchange='', routing_key='task_queue'): print(" [X] %s UTC %r %r" % (str(datetime.datetime.utcnow()), str(msg['id']), str(msg['file_path']))) -def run_proc(args, output=subprocess.DEVNULL): - proc = subprocess.Popen(args, stdout=output) - return proc.communicate() - def run_mercury(path): with tempfile.TemporaryDirectory() as tempdir: mercury = shutil.which('pmercury') mercury_output = os.path.join(tempdir, 'mercury_output.txt') args = [mercury, '-awxg', '-r', path, '-f', mercury_output] - run_proc(args) + network_tools_lib.run_proc(args) with open(mercury_output, 'r') as f: return f.read() diff --git a/network_tools_lib/network_tools_lib.py b/network_tools_lib/network_tools_lib.py index 3d536e00..d398a496 100644 --- a/network_tools_lib/network_tools_lib.py +++ b/network_tools_lib/network_tools_lib.py @@ -1,4 +1,5 @@ import os +import subprocess def mod_path(filename, file_for_dir=None): @@ -11,3 +12,9 @@ def get_version(): ver_path = os.path.join(mod_path('VERSION', __file__)) with open(ver_path, 'r') as f: return f.read().strip() + + +def run_proc(args, output=subprocess.DEVNULL): + with subprocess.Popen(args, stdout=output) as proc: + proc_output = proc.communicate() + return proc_output diff --git a/p0f/app.py b/p0f/app.py index 39d737a1..afd443de 100644 --- a/p0f/app.py +++ b/p0f/app.py @@ -3,7 +3,6 @@ import logging import os import shutil -import subprocess import sys import tempfile @@ -13,10 +12,6 @@ VERSION = network_tools_lib.get_version() -def run_proc(args, output=subprocess.DEVNULL): - proc = subprocess.Popen(args, stdout=output) - return proc.communicate() - def run_p0f(path): with tempfile.TemporaryDirectory() as tempdir: p0f = shutil.which('p0f') @@ -25,7 +20,7 @@ def run_p0f(path): p0f = '/usr/sbin/p0f' p0f_output = os.path.join(tempdir, 'p0f_output.txt') args = [p0f, '-r', path, '-o', p0f_output] - run_proc(args) + network_tools_lib.run_proc(args) with open(p0f_output, 'r') as f: return f.read() diff --git a/p0f/test_app.py b/p0f/test_app.py index d1f0c9ac..2cbf50e0 100644 --- a/p0f/test_app.py +++ b/p0f/test_app.py @@ -5,6 +5,7 @@ @author: Charlie Lewis """ import os +import shutil import sys from .app import VERSION @@ -14,11 +15,17 @@ from .app import run_tshark from .app import run_p0f from .app import build_result_json +import network_tools_lib TEST_LO_CAP = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'test_lo.cap') +def test_runproc(): + # returned filehandles both none as process has exited. + assert network_tools_lib.run_proc([shutil.which('ls')]) == (None, None) + + def test_ispcap(): assert ispcap('afile.pcap') # nosec assert not ispcap('notapcap.txt') # nosec