diff --git a/src/execution/logging.py b/src/execution/logging.py index f5a92b82..7d00f2c6 100644 --- a/src/execution/logging.py +++ b/src/execution/logging.py @@ -350,7 +350,7 @@ def create_filename(self, execution_id, all_audit_names, script_name, start_time date_string = ms_to_datetime(start_time).strftime(self._date_format) - username = get_first_existing(all_audit_names, audit_utils.AUTH_USERNAME, audit_utils.PROXIED_USERNAME) + username = audit_utils.get_audit_username(all_audit_names) mapping = { 'ID': execution_id, diff --git a/src/features/file_download_feature.py b/src/features/file_download_feature.py index 5548032c..6e6f6e85 100644 --- a/src/features/file_download_feature.py +++ b/src/features/file_download_feature.py @@ -7,8 +7,9 @@ import utils.os_utils as os_utils import utils.string_utils as string_utils from execution.execution_service import ExecutionService -from model.model_helper import is_empty, fill_parameter_values +from model.model_helper import is_empty, fill_parameter_values, replace_auth_vars from react.observable import read_until_closed +from utils import audit_utils from utils.file_utils import create_unique_filename INLINE_IMAGE_TYPE = 'inline-image' @@ -99,10 +100,17 @@ def _get_paths(self, execution_id, predicate): paths = [p for p in paths if p] parameter_values = self.execution_service.get_user_parameter_values(execution_id) - return substitute_parameter_values( + all_audit_names = self.execution_service.get_all_audit_names(execution_id) + + audit_name = audit_utils.get_audit_name(all_audit_names) + username = audit_utils.get_audit_username(all_audit_names) + + return substitute_variable_values( config.parameters, paths, - parameter_values) + parameter_values, + audit_name, + username) @staticmethod def _is_post_finish_path(file): @@ -228,10 +236,11 @@ def _add_inline_image(self, original_path, download_path): LOGGER.error('Failed to notify image listener') -def substitute_parameter_values(parameter_configs, output_files, values): +def substitute_variable_values(parameter_configs, output_files, values, audit_name, username): output_file_parsed = [] for _, output_file in enumerate(output_files): substituted_file = fill_parameter_values(parameter_configs, output_file, values) + substituted_file = replace_auth_vars(substituted_file, username, audit_name) output_file_parsed.append(substituted_file) return output_file_parsed diff --git a/src/tests/audit_utils_test.py b/src/tests/audit_utils_test.py index 04cf9dfb..e6ec3ce0 100644 --- a/src/tests/audit_utils_test.py +++ b/src/tests/audit_utils_test.py @@ -3,15 +3,17 @@ from tests.test_utils import mock_object from utils import audit_utils, os_utils -from auth.identification import AuthBasedIdentification +from utils.audit_utils import get_audit_username + def mock_request_handler(ip=None, proxy_username=None, auth_username=None, proxied_ip=None): handler_mock = mock_object() handler_mock.application = mock_object() - + handler_mock.application.identification = mock_object() handler_mock.application.identification.identify_for_audit = lambda x: auth_username + handler_mock.request = mock_object() handler_mock.request.headers = {} if proxy_username: @@ -128,3 +130,26 @@ def test_full_audit_localhost(self): 'auth_username': 'ldap_user', 'proxied_username': 'basic_username'}, names) + + +class TestGetAuditUsername(unittest.TestCase): + def test_auth_username(self): + username = get_audit_username({'auth_username': 'user_X', 'ip': '123'}) + self.assertEqual('user_X', username) + + def test_proxied_username(self): + username = get_audit_username({'proxied_username': 'Its me', 'hostname': '123'}) + self.assertEqual('Its me', username) + + def test_no_username(self): + username = get_audit_username({ + 'proxied_hostname': 'my host', + 'hostname': 'localhost', + 'ip': '123.456', + 'proxied_ip': '987'}) + self.assertIsNone(username) + + def test_auth_username_and_proxied_username(self): + username = get_audit_username({'proxied_username': 'Its me', 'auth_username': 'user_X'}) + self.assertEqual('user_X', username) + diff --git a/src/tests/file_download_feature_test.py b/src/tests/file_download_feature_test.py index eaeacf25..67b49acd 100644 --- a/src/tests/file_download_feature_test.py +++ b/src/tests/file_download_feature_test.py @@ -197,17 +197,22 @@ def tearDown(self): class TestParametersSubstitute(unittest.TestCase): def test_no_parameters(self): - files = file_download_feature.substitute_parameter_values([], ['/home/user/test.txt'], []) + files = file_download_feature.substitute_variable_values( + [], ['/home/user/test.txt'], [], + '127.0.0.1', + 'user_X') self.assertEqual(files, ['/home/user/test.txt']) def test_single_replace(self): parameter = create_parameter_model('param1') - files = file_download_feature.substitute_parameter_values( + files = file_download_feature.substitute_variable_values( [parameter], ['/home/user/${param1}.txt'], - {'param1': 'val1'}) + {'param1': 'val1'}, + '127.0.0.1', + 'user_X') self.assertEqual(files, ['/home/user/val1.txt']) @@ -216,10 +221,12 @@ def test_two_replaces(self): parameters.append(create_parameter_model('param1', all_parameters=parameters)) parameters.append(create_parameter_model('param2', all_parameters=parameters)) - files = file_download_feature.substitute_parameter_values( + files = file_download_feature.substitute_variable_values( parameters, ['/home/${param2}/${param1}.txt'], - {'param1': 'val1', 'param2': 'val2'}) + {'param1': 'val1', 'param2': 'val2'}, + '127.0.0.1', + 'user_X') self.assertEqual(files, ['/home/val2/val1.txt']) @@ -228,43 +235,87 @@ def test_two_replaces_in_two_files(self): parameters.append(create_parameter_model('param1', all_parameters=parameters)) parameters.append(create_parameter_model('param2', all_parameters=parameters)) - files = file_download_feature.substitute_parameter_values( + files = file_download_feature.substitute_variable_values( parameters, ['/home/${param2}/${param1}.txt', '/tmp/${param2}.txt', '/${param1}'], - {'param1': 'val1', 'param2': 'val2'}) + {'param1': 'val1', 'param2': 'val2'}, + '127.0.0.1', + 'user_X') self.assertEqual(files, ['/home/val2/val1.txt', '/tmp/val2.txt', '/val1']) def test_no_pattern_match(self): param1 = create_parameter_model('param1') - files = file_download_feature.substitute_parameter_values( + files = file_download_feature.substitute_variable_values( [param1], ['/home/user/${paramX}.txt'], - {'param1': 'val1'}) + {'param1': 'val1'}, + '127.0.0.1', + 'user_X') self.assertEqual(files, ['/home/user/${paramX}.txt']) def test_skip_secure_replace(self): param1 = create_parameter_model('param1', secure=True) - files = file_download_feature.substitute_parameter_values( + files = file_download_feature.substitute_variable_values( [param1], ['/home/user/${param1}.txt'], - {'param1': 'val1'}) + {'param1': 'val1'}, + '127.0.0.1', + 'user_X') self.assertEqual(files, ['/home/user/${param1}.txt']) def test_skip_flag_replace(self): param1 = create_parameter_model('param1', no_value=True) - files = file_download_feature.substitute_parameter_values( + files = file_download_feature.substitute_variable_values( [param1], ['/home/user/${param1}.txt'], - {'param1': 'val1'}) + {'param1': 'val1'}, + '127.0.0.1', + 'user_X') self.assertEqual(files, ['/home/user/${param1}.txt']) + def test_replace_audit_name(self): + param1 = create_parameter_model('param1', no_value=True) + + files = file_download_feature.substitute_variable_values( + [param1], + ['/home/user/${auth.audit_name}.txt'], + {'param1': 'val1'}, + '127.0.0.1', + 'user_X') + + self.assertEqual(files, ['/home/user/127.0.0.1.txt']) + + def test_replace_username(self): + param1 = create_parameter_model('param1', no_value=True) + + files = file_download_feature.substitute_variable_values( + [param1], + ['/home/user/${auth.username}.txt'], + {'param1': 'val1'}, + '127.0.0.1', + 'user_X') + + self.assertEqual(files, ['/home/user/user_X.txt']) + + def test_replace_username_and_param(self): + param1 = create_parameter_model('param1') + + files = file_download_feature.substitute_variable_values( + [param1], + ['/home/${auth.username}/${param1}.txt'], + {'param1': 'val1'}, + '127.0.0.1', + 'user_X') + + self.assertEqual(files, ['/home/user_X/val1.txt']) + class FileDownloadFeatureTest(unittest.TestCase): diff --git a/src/utils/audit_utils.py b/src/utils/audit_utils.py index 724f2d9b..1e06f86b 100644 --- a/src/utils/audit_utils.py +++ b/src/utils/audit_utils.py @@ -3,6 +3,7 @@ import socket import sys +from utils.collection_utils import get_first_existing from utils.tornado_utils import get_proxied_ip HOSTNAME = 'hostname' @@ -82,3 +83,7 @@ def find_basic_auth_username(request_handler): username = credentials.split(':')[0] return username + + +def get_audit_username(all_audit_names): + return get_first_existing(all_audit_names, AUTH_USERNAME, PROXIED_USERNAME)