PurpleDome/tests/test_machinecontrol.py

import unittest
import os
from app.machinecontrol import Machine
from app.exceptions import ConfigurationError
from app.config import MachineConfig
from unittest.mock import patch

# https://docs.python.org/3/library/unittest.html


class TestMachineControl(unittest.TestCase):

    def test_get_os_linux_machine(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3"})
        self.assertEqual(m.get_os(), "linux")

    def test_get_os_linux_machine_with_config_class(self):
        mc = MachineConfig({"root": "systems/attacker1",
                            "os": "linux",
                            "vm_controller": {
                                "type": "vagrant",
                                "vagrantfilepath": "systems",
                            },
                            "vm_name": "target3"})
        m = Machine(mc)
        self.assertEqual(m.get_os(), "linux")

    def test_get_os_missing(self):
        with self.assertRaises(ConfigurationError):
            Machine({"root": "systems/attacker1",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3"
                     })

    def test_get_os_not_supported(self):
        with self.assertRaises(ConfigurationError):
            Machine({"root": "systems/attacker1",
                     "os": "nintendo_switch",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3"})

    def test_get_paw_good(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "paw": "testme",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3"})
        self.assertEqual(m.get_paw(), "testme")

    def test_get_paw_missing(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3"
                     })
        self.assertEqual(m.get_paw(), None)

    def test_get_group_good(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "group": "testme",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3"})
        self.assertEqual(m.get_group(), "testme")

    def test_get_group_missing(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3"
                     })
        self.assertEqual(m.get_group(), None)

    def test_vagrantfilepath_missing(self):
        with self.assertRaises(ConfigurationError):
            Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                     },
                     "vm_name": "target3"
                     })

    def test_vagrantfile_missing(self):
        with self.assertRaises(ConfigurationError):
            Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "non_existing",
                     },
                     "vm_name": "target3"
                     })

    def test_vagrantfile_existing(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3"
                     })
        self.assertIsNotNone(m)

    def test_name_missing(self):
        with self.assertRaises(ConfigurationError):
            Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     })

    # test: auto generated, dir missing
    def test_auto_generated_machinepath_with_path_missing(self):
        with self.assertRaises(ConfigurationError):
            Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "missing"
                     })

    # test manual config, dir missing
    def test_configured_machinepath_with_path_missing(self):
        with self.assertRaises(ConfigurationError):
            Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3",
                     "machinepath": "missing"
                     })

    # test auto generated, dir there (external/internal dirs must work !)
    def test_auto_generated_machinepath_with_good_config(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3"
                     })
        vagrantfilepath = os.path.abspath("systems")
        ext = os.path.join(vagrantfilepath, "target3")
        internal = os.path.join("/vagrant/", "target3")

        self.assertEqual(m.abs_machinepath_external, ext)
        self.assertEqual(m.abs_machinepath_internal, internal)

    # test: manual config, dir there (external/internal dirs must work !)
    def test_configured_machinepath_with_good_config(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "missing",
                     "machinepath": "target3"
                     })
        vagrantfilepath = os.path.abspath("systems")
        ext = os.path.join(vagrantfilepath, "target3")
        internal = os.path.join("/vagrant/", "target3")

        self.assertEqual(m.abs_machinepath_external, ext)
        self.assertEqual(m.abs_machinepath_internal, internal)

    # vm_controller missing
    def test_configured_vm_controller_missing(self):
        with self.assertRaises(ConfigurationError):
            Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_name": "missing",
                     "machinepath": "target3"
                     })

    # vm_controller wrong
    def test_configured_vm_controller_wrong_type(self):
        with self.assertRaises(ConfigurationError):
            Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "wrong_controller",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "missing",
                     "machinepath": "target3"
                     })

    # Create caldera start command and verify it
    def test_get_linux_caldera_start_cmd(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "linux",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3",
                     "group": "testgroup",
                     "paw": "testpaw"})
        m.set_caldera_server("http://www.test.test")
        with patch.object(m.vm_manager, "get_playground", return_value="/vagrant/target3"):
            cmd = m.create_start_caldera_client_cmd()
        self.assertEqual(cmd.strip(), "chmod +x caldera_agent.sh; nohup bash /vagrant/target3/caldera_agent.sh start &".strip())

    # Create caldera start command and verify it (windows)
    def test_get_windows_caldera_start_cmd(self):
        m = Machine({"root": "systems/attacker1",
                     "os": "windows",
                     "vm_controller": {
                         "type": "vagrant",
                         "vagrantfilepath": "systems",
                     },
                     "vm_name": "target3",
                     "group": "testgroup",
                     "paw": "testpaw",
                     "machinepath": "target2w"})
        m.set_caldera_server("www.test.test")
        cmd = m.create_start_caldera_client_cmd()
        self.maxDiff = None
        expected = """
caldera_agent.bat"""
        self.assertEqual(cmd.strip(), expected.strip())


if __name__ == '__main__':
    unittest.main()