island: Refactor get_file_descriptor_for_new_secure_file as contextmgr

get_file_descriptor_for_new_secure_file() has been refactored as a
contextmanager. Additionally, it has been renamed to
open_new_securely_permissioned_file(). The function can now be used
similarly to open().

Example:

with open_new_securely_permissioned_file(file_path, "wb") as f:
    f.write(data)
This commit is contained in:
Mike Salvatore 2021-06-19 09:12:04 -04:00
parent b39440e871
commit 51aa0d1564
3 changed files with 34 additions and 15 deletions

View File

@ -6,7 +6,7 @@ import os
from Crypto import Random # noqa: DUO133 # nosec: B413 from Crypto import Random # noqa: DUO133 # nosec: B413
from Crypto.Cipher import AES # noqa: DUO133 # nosec: B413 from Crypto.Cipher import AES # noqa: DUO133 # nosec: B413
from monkey_island.cc.server_utils.file_utils import get_file_descriptor_for_new_secure_file from monkey_island.cc.server_utils.file_utils import open_new_securely_permissioned_file
__author__ = "itay.mizeretz" __author__ = "itay.mizeretz"
@ -27,8 +27,7 @@ class Encryptor:
def _init_key(self, password_file_path: str): def _init_key(self, password_file_path: str):
self._cipher_key = Random.new().read(self._BLOCK_SIZE) self._cipher_key = Random.new().read(self._BLOCK_SIZE)
get_file_descriptor_for_new_secure_file(path=password_file_path) with open_new_securely_permissioned_file(password_file_path, "wb") as f:
with open(password_file_path, "wb") as f:
f.write(self._cipher_key) f.write(self._cipher_key)
def _load_existing_key(self, password_file): def _load_existing_key(self, password_file):

View File

@ -2,6 +2,8 @@ import logging
import os import os
import platform import platform
import stat import stat
from contextlib import contextmanager
from typing import Generator
LOG = logging.getLogger(__name__) LOG = logging.getLogger(__name__)
@ -54,11 +56,15 @@ def _create_secure_directory_windows(path: str):
raise ex raise ex
def get_file_descriptor_for_new_secure_file(path: str) -> int: @contextmanager
def open_new_securely_permissioned_file(path: str, mode: str = "w") -> Generator:
if is_windows_os(): if is_windows_os():
return _get_file_descriptor_for_new_secure_file_windows(path) fd = _get_file_descriptor_for_new_secure_file_windows(path)
else: else:
return _get_file_descriptor_for_new_secure_file_linux(path) fd = _get_file_descriptor_for_new_secure_file_linux(path)
with open(fd, mode) as f:
yield f
def _get_file_descriptor_for_new_secure_file_linux(path: str) -> int: def _get_file_descriptor_for_new_secure_file_linux(path: str) -> int:

View File

@ -6,8 +6,8 @@ import pytest
from monkey_island.cc.server_utils.file_utils import ( from monkey_island.cc.server_utils.file_utils import (
create_secure_directory, create_secure_directory,
expand_path, expand_path,
get_file_descriptor_for_new_secure_file,
is_windows_os, is_windows_os,
open_new_securely_permissioned_file,
) )
if is_windows_os(): if is_windows_os():
@ -98,22 +98,26 @@ def test_create_secure_directory__perm_windows(test_path):
assert ace_inheritance == ACE_INHERIT_OBJECT_AND_CONTAINER assert ace_inheritance == ACE_INHERIT_OBJECT_AND_CONTAINER
def test_get_file_descriptor_for_new_secure_file__already_exists(test_path): def test_open_new_securely_permissioned_file__already_exists(test_path):
os.close(os.open(test_path, os.O_CREAT, stat.S_IRWXU)) os.close(os.open(test_path, os.O_CREAT, stat.S_IRWXU))
assert os.path.isfile(test_path) assert os.path.isfile(test_path)
with pytest.raises(Exception): with pytest.raises(Exception):
get_file_descriptor_for_new_secure_file(test_path) with open_new_securely_permissioned_file(test_path):
pass
def test_get_file_descriptor_for_new_secure_file__no_parent_dir(test_path_nested): def test_open_new_securely_permissioned_file__no_parent_dir(test_path_nested):
with pytest.raises(Exception): with pytest.raises(Exception):
get_file_descriptor_for_new_secure_file(test_path_nested) with open_new_securely_permissioned_file(test_path_nested):
pass
@pytest.mark.skipif(is_windows_os(), reason="Tests Posix (not Windows) permissions.") @pytest.mark.skipif(is_windows_os(), reason="Tests Posix (not Windows) permissions.")
def test_get_file_descriptor_for_new_secure_file__perm_linux(test_path): def test_open_new_securely_permissioned_file__perm_linux(test_path):
os.close(get_file_descriptor_for_new_secure_file(test_path)) with open_new_securely_permissioned_file(test_path):
pass
st = os.stat(test_path) st = os.stat(test_path)
expected_mode = stat.S_IRUSR | stat.S_IWUSR expected_mode = stat.S_IRUSR | stat.S_IWUSR
@ -123,8 +127,9 @@ def test_get_file_descriptor_for_new_secure_file__perm_linux(test_path):
@pytest.mark.skipif(not is_windows_os(), reason="Tests Windows (not Posix) permissions.") @pytest.mark.skipif(not is_windows_os(), reason="Tests Windows (not Posix) permissions.")
def test_get_file_descriptor_for_new_secure_file__perm_windows(test_path): def test_open_new_securely_permissioned_file__perm_windows(test_path):
os.close(get_file_descriptor_for_new_secure_file(test_path)) with open_new_securely_permissioned_file(test_path):
pass
acl, user_sid = _get_acl_and_sid_from_path(test_path) acl, user_sid = _get_acl_and_sid_from_path(test_path)
@ -140,3 +145,12 @@ def test_get_file_descriptor_for_new_secure_file__perm_windows(test_path):
assert ace_sid == user_sid assert ace_sid == user_sid
assert ace_permissions == FULL_CONTROL and ace_access_mode == ACE_ACCESS_MODE_GRANT_ACCESS assert ace_permissions == FULL_CONTROL and ace_access_mode == ACE_ACCESS_MODE_GRANT_ACCESS
assert ace_inheritance == ACE_INHERIT_OBJECT_AND_CONTAINER assert ace_inheritance == ACE_INHERIT_OBJECT_AND_CONTAINER
def test_open_new_securely_permissioned_file__write(test_path):
TEST_STR = b"Hello World"
with open_new_securely_permissioned_file(test_path, "wb") as f:
f.write(TEST_STR)
with open(test_path, "rb") as f:
assert f.read() == TEST_STR