From 1e425531126f1ae016dbb84248447638e39dbf94 Mon Sep 17 00:00:00 2001 From: hakril Date: Mon, 12 Aug 2019 13:53:15 +0200 Subject: [PATCH] Added some test + added a com security fixture allowing allowed correct COM init accross all tests --- tests/conftest.py | 8 ++++++++ tests/test_bits.py | 27 +++++++++++++++++++++++++++ tests/test_evtlog.py | 25 +++++++++++++++++++++++++ tests/test_system.py | 26 +++++++++++++++++++++----- tests/test_wmi.py | 32 ++++++++++---------------------- 5 files changed, 91 insertions(+), 27 deletions(-) create mode 100644 tests/test_bits.py diff --git a/tests/conftest.py b/tests/conftest.py index 95f39e9..bea605f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -58,6 +58,14 @@ else: proc32_64_suspended = generate_pop_and_exit_fixtures([pop_proc_32, pop_proc_64], ids=["proc32", "proc64"], dwCreationFlags=gdef.CREATE_SUSPENDED) +@pytest.fixture(scope="session") +def init_com_security(): + # Init com security if not done + try: + windows.com.init() + return windows.com.initsecurity() + except WindowsError: + pass class HandleDebugger(object): diff --git a/tests/test_bits.py b/tests/test_bits.py new file mode 100644 index 0000000..e6580b7 --- /dev/null +++ b/tests/test_bits.py @@ -0,0 +1,27 @@ +import pytest +import os.path + +import windows +import windows.generated_def as gdef + +pytestmark = pytest.mark.usefixtures("init_com_security") + +SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__)) + +@pytest.fixture +def bitsjob(): + newjob = windows.system.bits.create("PFW_TEST_BITSJOB", gdef.BG_JOB_TYPE_DOWNLOAD) + yield newjob + newjob.Cancel() + +def test_job_state(bitsjob): + # Check state returns the value and not enum struct + assert bitsjob.state == gdef.BG_JOB_STATE_SUSPENDED + # Enum value should be castable to int + assert int(bitsjob.state) == gdef.BG_JOB_STATE_SUSPENDED + +def test_job_multi_files(bitsjob): + bitsjob.AddFile("https://example.com/REMOTE_FILE_1", os.path.join(SCRIPT_PATH, "local1")) + bitsjob.AddFile("https://example.com/REMOTE_FILE_2", os.path.join(SCRIPT_PATH, "local2")) + files = bitsjob.files + assert len(files) == 2 diff --git a/tests/test_evtlog.py b/tests/test_evtlog.py index bdc7fd3..1eb1ab7 100644 --- a/tests/test_evtlog.py +++ b/tests/test_evtlog.py @@ -69,4 +69,29 @@ def test_new_event(): assert any(evt.pid == p.pid for evt in new_events) +def test_event_close(): + chan = windows.system.event_log["System"] + start_usage = windows.current_process.memory_info.PrivateUsage + count_max = 0x10000 + count = 0 + while count < count_max: + for i,e in enumerate(chan.query()): + count += 1 + post_usage = windows.current_process.memory_info.PrivateUsage + memory_usage_in_mo = (post_usage - start_usage) / 1024 / 1024 + memory_usage_in_ko = (post_usage - start_usage) / 1024 + # With auto-evtclose of evt there should not be too much memory used when + # Variable are not accessible anymore + assert memory_usage_in_mo == 0 +def test_evthandle_close(): + start_usage = windows.current_process.memory_info.PrivateUsage + for i in range(0x2000): + chan = windows.system.event_log["System"] + query = chan.query() + config = chan.config # Config is an EVT_HANDLE + pubm = config.publisher.metadata + # windows.winproxy.EvtClose(chan) + post_usage = windows.current_process.memory_info.PrivateUsage + memory_usage_in_mo = (post_usage - start_usage) / 1024 / 1024 + assert memory_usage_in_mo == 0 diff --git a/tests/test_system.py b/tests/test_system.py index 0d9ee7d..60c8a98 100644 --- a/tests/test_system.py +++ b/tests/test_system.py @@ -11,6 +11,16 @@ class TestSystemWithCheckGarbage(object): def test_version_name(self): return windows.system.version_name + def test_version_product_type(self): + return windows.system.product_type + + def test_version_edition(self): + return windows.system.edition + + def test_version_windir(self): + return windows.system.windir + + def test_computer_name(self): return windows.system.computer_name @@ -26,11 +36,17 @@ class TestSystemWithCheckGarbage(object): def test_handles(self): return windows.system.handles - def test_handle_process(self): - handle_with_process = [h for h in windows.system.handles if h.dwProcessId] - handle = handle_with_process[-1] - proc = handle.process - assert proc.pid == handle.dwProcessId + def test_bitness(self): + return windows.system.bitness + + def test_evtlog(self): + return windows.system.event_log + + def test_task_scheduler(self): + return windows.system.task_scheduler + + def test_task_object_manager(self): + return windows.system.object_manager def test_system_modules_ntosk(self): assert windows.system.modules[0].name.endswith("ntoskrnl.exe") diff --git a/tests/test_wmi.py b/tests/test_wmi.py index 9a5e48d..4878bea 100644 --- a/tests/test_wmi.py +++ b/tests/test_wmi.py @@ -12,32 +12,20 @@ from pfwtest import * # does not allow to perform the request we want.. # So we try & do it ourself here. - pytestmark = pytest.mark.usefixtures("init_com_security") -@pytest.fixture(scope="module") -def init_com_security(): - # Init com security if not done - try: - return windows.com.initsecurity() - except WindowsError: - pass - - -wmimanager = windows.system.wmi - @pytest.mark.parametrize("name, expected_cls", [ ("root\\cimv2", "Win32_Process"), ("root\\subscription", "__EventFilter"), ]) def test_wmimanager_getnamespace(name, expected_cls): - namespace = wmimanager[name] + namespace = windows.system.wmi[name] assert namespace.name == name assert namespace.get_object(expected_cls) def test_wmimanager_subnamespaces(): - subnamespaces = wmimanager.get_subnamespaces("root") + subnamespaces = windows.system.wmi.get_subnamespaces("root") subnamespaces = [x.lower() for x in subnamespaces] assert "cimv2" in subnamespaces assert "security" in subnamespaces @@ -50,20 +38,20 @@ def test_wmimanager_subnamespaces(): ("root\\subscription", "select * from __EventFilter"), ]) def test_query_select(name, query): - namespace = wmimanager[name] + namespace = windows.system.wmi[name] x = namespace.query(query) assert x assert isinstance(x, list) def test_bad_query_raise(): - namespace = wmimanager["root\\cimv2"] + namespace = windows.system.wmi["root\\cimv2"] with pytest.raises(WindowsError) as e: x = namespace.query("BADSELECT QUERY BAD") assert (e.value.winerror & 0xffffffff) == gdef.WBEM_E_INVALID_QUERY def test_create_class_enum(): - namespace = wmimanager["root\\cimv2"] + namespace = windows.system.wmi["root\\cimv2"] enum = namespace.create_class_enum(None) assert enum classes = list(enum) @@ -77,7 +65,7 @@ def test_create_class_enum(): ("root\\subscription", "__EventFilter"), ]) def test_get_object(name, cls): - namespace = wmimanager[name] + namespace = windows.system.wmi[name] assert namespace.name == name obj = namespace.get_object(cls) assert obj["__CLASS"] == cls @@ -88,7 +76,7 @@ def test_get_object(name, cls): @pytest.mark.parametrize("cmdline", [r"c:\windows\notepad.exe trolol.exe"]) def test_exec_method_Win32_Process_create(cmdline): - namespace = wmimanager["root\\cimv2"] + namespace = windows.system.wmi["root\\cimv2"] win32_process_cls = namespace.get_object("Win32_Process") inparam = win32_process_cls.get_method("Create").inparam.spawn_instance() inparam["CommandLine"] = cmdline @@ -103,7 +91,7 @@ def test_exec_method_Win32_Process_create(cmdline): ## Test enum def test_enumeration_iteration_no_timeout(): - namespace = wmimanager["root\\cimv2"] + namespace = windows.system.wmi["root\\cimv2"] processes = namespace.exec_query("select * from Win32_Process").all() assert isinstance(processes, list) assert processes @@ -115,7 +103,7 @@ def test_enumeration_iteration_no_timeout(): assert proc["__CLASS"].lower() == "win32_process" def test_enumeration_iteration_timeout(): - namespace = wmimanager["root\\cimv2"] + namespace = windows.system.wmi["root\\cimv2"] timegen = namespace.exec_query("select * from Win32_Process").iter_timeout(0) # Iter on Win32_Process should not be immediat # so itering on timegen should trigger a timeout @@ -128,7 +116,7 @@ def test_enumeration_iteration_timeout(): def wmi_cls(): # Test expect the cls to have a "Name" attribute & "Create" method # Maybe doing something more generic - namespace = wmimanager["root\\cimv2"] + namespace = windows.system.wmi["root\\cimv2"] yield namespace.get_object("Win32_Process") def test_wmiobject_spawn(wmi_cls):