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()