From 09739268fa58af4f6b1bea0e89cde49ab25c05af Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Mon, 22 Aug 2022 14:07:09 -0400 Subject: [PATCH] Agent: Add get_machine_id() --- .../utils/{agent_id.py => ids.py} | 7 ++++++- .../infection_monkey/utils/test_agent_id.py | 10 ---------- .../infection_monkey/utils/test_ids.py | 17 +++++++++++++++++ 3 files changed, 23 insertions(+), 11 deletions(-) rename monkey/infection_monkey/utils/{agent_id.py => ids.py} (70%) delete mode 100644 monkey/tests/unit_tests/infection_monkey/utils/test_agent_id.py create mode 100644 monkey/tests/unit_tests/infection_monkey/utils/test_ids.py diff --git a/monkey/infection_monkey/utils/agent_id.py b/monkey/infection_monkey/utils/ids.py similarity index 70% rename from monkey/infection_monkey/utils/agent_id.py rename to monkey/infection_monkey/utils/ids.py index 0a889a114..04f9aea29 100644 --- a/monkey/infection_monkey/utils/agent_id.py +++ b/monkey/infection_monkey/utils/ids.py @@ -1,4 +1,4 @@ -from uuid import UUID, uuid4 +from uuid import UUID, getnode, uuid4 def get_agent_id() -> UUID: @@ -15,3 +15,8 @@ def get_agent_id() -> UUID: get_agent_id._id = None + + +def get_machine_id() -> int: + """Get an integer that uniquely defines the machine the agent is running on""" + return getnode() diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_agent_id.py b/monkey/tests/unit_tests/infection_monkey/utils/test_agent_id.py deleted file mode 100644 index c4a95b73c..000000000 --- a/monkey/tests/unit_tests/infection_monkey/utils/test_agent_id.py +++ /dev/null @@ -1,10 +0,0 @@ -from uuid import UUID - -from infection_monkey.utils.agent_id import get_agent_id - - -def test_get_agent_id(): - agent_id = get_agent_id() - - assert isinstance(agent_id, UUID) - assert agent_id == get_agent_id() diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_ids.py b/monkey/tests/unit_tests/infection_monkey/utils/test_ids.py new file mode 100644 index 000000000..28fa9a64e --- /dev/null +++ b/monkey/tests/unit_tests/infection_monkey/utils/test_ids.py @@ -0,0 +1,17 @@ +from uuid import UUID + +from infection_monkey.utils.ids import get_agent_id, get_machine_id + + +def test_get_agent_id(): + agent_id = get_agent_id() + + assert isinstance(agent_id, UUID) + assert agent_id == get_agent_id() + + +def test_get_machine_id(): + machine_id = get_machine_id() + + assert isinstance(machine_id, int) + assert machine_id == get_machine_id()