Docker capture and refactor common capture code

Fix https://github.com/GNS3/gns3-gui/issues/891

gns3server/handlers/api/docker_handler.py

@@ -15,6 +15,7 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import os
 from aiohttp.web import HTTPConflict
 
 from ...web.route import Route
@@ -26,6 +27,7 @@ from ...schemas.docker import (
     DOCKER_UPDATE_SCHEMA,
     DOCKER_LIST_IMAGES_SCHEMA
 )
+from ...schemas.vm import VM_CAPTURE_SCHEMA
 from ...schemas.nio import NIO_SCHEMA
 
 
@@ -276,3 +278,60 @@ class DockerHandler:
         vm.environment = request.json.get("environment", vm.environment)
         yield from vm.update()
         response.json(vm)
+
+    @Route.post(
+        r"/projects/{project_id}/docker/vms/{vm_id}/adapters/{adapter_number:\d+}/ports/{port_number:\d+}/start_capture",
+        parameters={
+            "project_id": "UUID for the project",
+            "vm_id": "UUID for the instance",
+            "adapter_number": "Adapter to start a packet capture",
+            "port_number": "Port on the adapter"
+        },
+        status_codes={
+            200: "Capture started",
+            400: "Invalid request",
+            404: "Instance doesn't exist",
+            409: "VM not started"
+        },
+        description="Start a packet capture on a IOU VM instance",
+        input=VM_CAPTURE_SCHEMA)
+    def start_capture(request, response):
+
+        iou_manager = Docker.instance()
+        vm = iou_manager.get_vm(request.match_info["vm_id"], project_id=request.match_info["project_id"])
+        adapter_number = int(request.match_info["adapter_number"])
+        port_number = int(request.match_info["port_number"])
+        pcap_file_path = os.path.join(vm.project.capture_working_directory(), request.json["capture_file_name"])
+
+        if not vm.is_running():
+            raise HTTPConflict(text="Cannot capture traffic on a non started VM")
+        yield from vm.start_capture(adapter_number, pcap_file_path)
+        response.json({"pcap_file_path": str(pcap_file_path)})
+
+    @Route.post(
+        r"/projects/{project_id}/docker/vms/{vm_id}/adapters/{adapter_number:\d+}/ports/{port_number:\d+}/stop_capture",
+        parameters={
+            "project_id": "UUID for the project",
+            "vm_id": "UUID for the instance",
+            "adapter_number": "Adapter to stop a packet capture",
+            "port_number": "Port on the adapter (always 0)"
+        },
+        status_codes={
+            204: "Capture stopped",
+            400: "Invalid request",
+            404: "Instance doesn't exist",
+            409: "VM not started"
+        },
+        description="Stop a packet capture on a IOU VM instance")
+    def stop_capture(request, response):
+
+        iou_manager = Docker.instance()
+        vm = iou_manager.get_vm(request.match_info["vm_id"], project_id=request.match_info["project_id"])
+
+        if not vm.is_running():
+            raise HTTPConflict(text="Cannot capture traffic on a non started VM")
+
+        adapter_number = int(request.match_info["adapter_number"])
+        port_number = int(request.match_info["port_number"])
+        yield from vm.stop_capture(adapter_number, port_number)
+        response.set_status(204)

gns3server/handlers/api/dynamips_device_handler.py

@@ -20,9 +20,9 @@ import asyncio
 from ...web.route import Route
 from ...schemas.dynamips_device import DEVICE_CREATE_SCHEMA
 from ...schemas.dynamips_device import DEVICE_UPDATE_SCHEMA
-from ...schemas.dynamips_device import DEVICE_CAPTURE_SCHEMA
 from ...schemas.dynamips_device import DEVICE_OBJECT_SCHEMA
 from ...schemas.dynamips_device import DEVICE_NIO_SCHEMA
+from ...schemas.vm import VM_CAPTURE_SCHEMA
 from ...modules.dynamips import Dynamips
 
 
@@ -198,7 +198,7 @@ class DynamipsDeviceHandler:
             404: "Instance doesn't exist"
         },
         description="Start a packet capture on a Dynamips device instance",
-        input=DEVICE_CAPTURE_SCHEMA)
+        input=VM_CAPTURE_SCHEMA)
     def start_capture(request, response):
 
         dynamips_manager = Dynamips.instance()

gns3server/handlers/api/dynamips_vm_handler.py

@@ -24,9 +24,9 @@ from ...schemas.nio import NIO_SCHEMA
 from ...schemas.vm import VM_LIST_IMAGES_SCHEMA
 from ...schemas.dynamips_vm import VM_CREATE_SCHEMA
 from ...schemas.dynamips_vm import VM_UPDATE_SCHEMA
-from ...schemas.dynamips_vm import VM_CAPTURE_SCHEMA
 from ...schemas.dynamips_vm import VM_OBJECT_SCHEMA
 from ...schemas.dynamips_vm import VM_CONFIGS_SCHEMA
+from ...schemas.vm import VM_CAPTURE_SCHEMA
 from ...modules.dynamips import Dynamips
 from ...modules.dynamips.dynamips_error import DynamipsError
 from ...modules.project_manager import ProjectManager

gns3server/handlers/api/iou_handler.py

@@ -24,9 +24,9 @@ from ...schemas.iou import IOU_CREATE_SCHEMA
 from ...schemas.iou import IOU_START_SCHEMA
 from ...schemas.iou import IOU_UPDATE_SCHEMA
 from ...schemas.iou import IOU_OBJECT_SCHEMA
-from ...schemas.iou import IOU_CAPTURE_SCHEMA
 from ...schemas.iou import IOU_CONFIGS_SCHEMA
 from ...schemas.vm import VM_LIST_IMAGES_SCHEMA
+from ...schemas.vm import VM_CAPTURE_SCHEMA
 from ...modules.iou import IOU
 
 
@@ -274,7 +274,7 @@ class IOUHandler:
             409: "VM not started"
         },
         description="Start a packet capture on a IOU VM instance",
-        input=IOU_CAPTURE_SCHEMA)
+        input=VM_CAPTURE_SCHEMA)
     def start_capture(request, response):
 
         iou_manager = IOU.instance()

gns3server/handlers/api/virtualbox_handler.py

@@ -22,8 +22,8 @@ from ...web.route import Route
 from ...schemas.nio import NIO_SCHEMA
 from ...schemas.virtualbox import VBOX_CREATE_SCHEMA
 from ...schemas.virtualbox import VBOX_UPDATE_SCHEMA
-from ...schemas.virtualbox import VBOX_CAPTURE_SCHEMA
 from ...schemas.virtualbox import VBOX_OBJECT_SCHEMA
+from ...schemas.vm import VM_CAPTURE_SCHEMA
 from ...modules.virtualbox import VirtualBox
 from ...modules.project_manager import ProjectManager
 
@@ -342,7 +342,7 @@ class VirtualBoxHandler:
             404: "Instance doesn't exist"
         },
         description="Start a packet capture on a VirtualBox VM instance",
-        input=VBOX_CAPTURE_SCHEMA)
+        input=VM_CAPTURE_SCHEMA)
     def start_capture(request, response):
 
         vbox_manager = VirtualBox.instance()

gns3server/handlers/api/vmware_handler.py

@@ -22,7 +22,7 @@ from ...web.route import Route
 from ...schemas.vmware import VMWARE_CREATE_SCHEMA
 from ...schemas.vmware import VMWARE_UPDATE_SCHEMA
 from ...schemas.vmware import VMWARE_OBJECT_SCHEMA
-from ...schemas.vmware import VMWARE_CAPTURE_SCHEMA
+from ...schemas.vm import VM_CAPTURE_SCHEMA
 from ...schemas.nio import NIO_SCHEMA
 from ...modules.vmware import VMware
 from ...modules.project_manager import ProjectManager
@@ -314,7 +314,7 @@ class VMwareHandler:
             404: "Instance doesn't exist",
         },
         description="Start a packet capture on a VMware VM instance",
-        input=VMWARE_CAPTURE_SCHEMA)
+        input=VM_CAPTURE_SCHEMA)
     def start_capture(request, response):
 
         vmware_manager = VMware.instance()

gns3server/modules/docker/docker_vm.py

@@ -501,3 +501,89 @@ class DockerVM(BaseVM):
             except ValueError:  # Partial JSON
                 pass
         self.project.emit("log.info", {"message": "Success pulling image {}".format(self._image)})
+
+    @asyncio.coroutine
+    def _start_ubridge_capture(self, adapter_number, output_file):
+        """
+        Start a packet capture in uBridge.
+
+        :param adapter_number: adapter number
+        :param output_file: PCAP destination file for the capture
+        """
+
+        adapter = "bridge{}".format(adapter_number)
+        if not self._ubridge_hypervisor:
+            raise VMwareError("Cannot start the packet capture: uBridge is not running")
+        yield from self._ubridge_hypervisor.send('bridge start_capture {name} "{output_file}"'.format(name=adapter, output_file=output_file))
+
+    @asyncio.coroutine
+    def _stop_ubridge_capture(self, adapter_number):
+        """
+        Stop a packet capture in uBridge.
+
+        :param adapter_number: adapter number
+        """
+
+        adapter = "bridge{}".format(adapter_number)
+        if not self._ubridge_hypervisor:
+            raise VMwareError("Cannot stop the packet capture: uBridge is not running")
+        yield from self._ubridge_hypervisor.send("bridge stop_capture {name}".format(name=adapter))
+
+    @asyncio.coroutine
+    def start_capture(self, adapter_number, output_file):
+        """
+        Starts a packet capture.
+
+        :param adapter_number: adapter number
+        :param output_file: PCAP destination file for the capture
+        """
+
+        try:
+            adapter = self._ethernet_adapters[adapter_number]
+        except KeyError:
+            raise DockerError("Adapter {adapter_number} doesn't exist on Docker VM '{name}'".format(name=self.name,
+                                                                                                    adapter_number=adapter_number))
+
+        nio = adapter.get_nio(0)
+
+        if not nio:
+            raise DockerError("Adapter {} is not connected".format(adapter_number))
+
+        if nio.capturing:
+            raise DockerError("Packet capture is already activated on adapter {adapter_number}".format(adapter_number=adapter_number))
+
+        nio.startPacketCapture(output_file)
+
+        if self.status == "started":
+            yield from self._start_ubridge_capture(adapter_number, output_file)
+
+        log.info("Docker VM '{name}' [{id}]: starting packet capture on adapter {adapter_number}".format(name=self.name,
+                                                                                                         id=self.id,
+                                                                                                         adapter_number=adapter_number))
+
+    def stop_capture(self, adapter_number):
+        """
+        Stops a packet capture.
+
+        :param adapter_number: adapter number
+        """
+
+        try:
+            adapter = self._ethernet_adapters[adapter_number]
+        except KeyError:
+            raise DockerError("Adapter {adapter_number} doesn't exist on Docker VM '{name}'".format(name=self.name,
+                                                                                                    adapter_number=adapter_number))
+
+        nio = adapter.get_nio(0)
+
+        if not nio:
+            raise DockerError("Adapter {} is not connected".format(adapter_number))
+
+        nio.stopPacketCapture()
+
+        if self.status == "started":
+            yield from self._stop_ubridge_capture(adapter_number)
+
+        log.info("Docker VM '{name}' [{id}]: stopping packet capture on adapter {adapter_number}".format(name=self.name,
+                                                                                                         id=self.id,
+                                                                                                         adapter_number=adapter_number))

gns3server/modules/vmware/vmware_vm.py

@@ -976,8 +976,8 @@ class VMwareVM(BaseVM):
         try:
             adapter = self._ethernet_adapters[adapter_number]
         except KeyError:
-            raise VMwareError("Adapter {adapter_number} doesn't exist on VirtualBox VM '{name}'".format(name=self.name,
-                                                                                                        adapter_number=adapter_number))
+            raise VMwareError("Adapter {adapter_number} doesn't exist on VMware VM '{name}'".format(name=self.name,
+                                                                                                    adapter_number=adapter_number))
 
         nio = adapter.get_nio(0)
 

gns3server/schemas/dynamips_device.py

@@ -346,23 +346,3 @@ DEVICE_NIO_SCHEMA = {
     "additionalProperties": False,
     "required": ["nio"]
 }
-
-DEVICE_CAPTURE_SCHEMA = {
-    "$schema": "http://json-schema.org/draft-04/schema#",
-    "description": "Request validation to start a packet capture on an Device instance port",
-    "type": "object",
-    "properties": {
-        "capture_file_name": {
-            "description": "Capture file name",
-            "type": "string",
-            "minLength": 1,
-        },
-        "data_link_type": {
-            "description": "PCAP data link type",
-            "type": "string",
-            "minLength": 1,
-        },
-    },
-    "additionalProperties": False,
-    "required": ["capture_file_name", "data_link_type"]
-}

gns3server/schemas/dynamips_vm.py

@@ -491,25 +491,6 @@ VM_UPDATE_SCHEMA = {
     "additionalProperties": False,
 }
 
-VM_CAPTURE_SCHEMA = {
-    "$schema": "http://json-schema.org/draft-04/schema#",
-    "description": "Request validation to start a packet capture on a Dynamips VM instance port",
-    "type": "object",
-    "properties": {
-        "capture_file_name": {
-            "description": "Capture file name",
-            "type": "string",
-            "minLength": 1,
-        },
-        "data_link_type": {
-            "description": "PCAP data link type",
-            "type": "string",
-            "minLength": 1,
-        },
-    },
-    "additionalProperties": False,
-    "required": ["capture_file_name", "data_link_type"]
-}
 
 VM_OBJECT_SCHEMA = {
     "$schema": "http://json-schema.org/draft-04/schema#",

gns3server/schemas/iou.py

@@ -266,25 +266,6 @@ IOU_OBJECT_SCHEMA = {
                  "command_line"]
 }
 
-IOU_CAPTURE_SCHEMA = {
-    "$schema": "http://json-schema.org/draft-04/schema#",
-    "description": "Request validation to start a packet capture on a IOU instance",
-    "type": "object",
-    "properties": {
-        "capture_file_name": {
-            "description": "Capture file name",
-            "type": "string",
-            "minLength": 1,
-        },
-        "data_link_type": {
-            "description": "PCAP data link type",
-            "type": "string",
-            "minLength": 1,
-        },
-    },
-    "additionalProperties": False,
-    "required": ["capture_file_name", "data_link_type"]
-}
 
 IOU_CONFIGS_SCHEMA = {
     "$schema": "http://json-schema.org/draft-04/schema#",

gns3server/schemas/virtualbox.py

@@ -147,20 +147,6 @@ VBOX_UPDATE_SCHEMA = {
     "additionalProperties": False,
 }
 
-VBOX_CAPTURE_SCHEMA = {
-    "$schema": "http://json-schema.org/draft-04/schema#",
-    "description": "Request validation to start a packet capture on a VirtualBox VM instance port",
-    "type": "object",
-    "properties": {
-        "capture_file_name": {
-            "description": "Capture file name",
-            "type": "string",
-            "minLength": 1,
-        },
-    },
-    "additionalProperties": False,
-    "required": ["capture_file_name"]
-}
 
 VBOX_OBJECT_SCHEMA = {
     "$schema": "http://json-schema.org/draft-04/schema#",

gns3server/schemas/vm.py

@@ -41,3 +41,24 @@ VM_LIST_IMAGES_SCHEMA = {
     ],
     "additionalProperties": False,
 }
+
+
+VM_CAPTURE_SCHEMA = {
+    "$schema": "http://json-schema.org/draft-04/schema#",
+    "description": "Request validation to start a packet capture on a port",
+    "type": "object",
+    "properties": {
+        "capture_file_name": {
+            "description": "Capture file name",
+            "type": "string",
+            "minLength": 1,
+        },
+        "data_link_type": {
+            "description": "PCAP data link type",
+            "type": "string",
+            "minLength": 1,
+        }
+    },
+    "additionalProperties": False,
+    "required": ["capture_file_name"]
+}

gns3server/schemas/vmware.py

@@ -140,20 +140,6 @@ VMWARE_UPDATE_SCHEMA = {
     "additionalProperties": False,
 }
 
-VMWARE_CAPTURE_SCHEMA = {
-    "$schema": "http://json-schema.org/draft-04/schema#",
-    "description": "Request validation to start a packet capture on a VMware VM instance port",
-    "type": "object",
-    "properties": {
-        "capture_file_name": {
-            "description": "Capture file name",
-            "type": "string",
-            "minLength": 1,
-        },
-    },
-    "additionalProperties": False,
-    "required": ["capture_file_name"]
-}
 
 VMWARE_OBJECT_SCHEMA = {
     "$schema": "http://json-schema.org/draft-04/schema#",

tests/handlers/api/test_docker.py

@@ -123,3 +123,52 @@ def test_docker_update(server, vm, tmpdir, free_console_port):
     assert response.json["console"] == free_console_port
     assert response.json["start_command"] == "yes"
     assert response.json["environment"] == "GNS3=1\nGNS4=0"
+
+
+def test_docker_start_capture(server, vm, tmpdir, project):
+
+    with patch("gns3server.modules.docker.docker_vm.DockerVM.is_running", return_value=True) as mock:
+        with asyncio_patch("gns3server.modules.docker.docker_vm.DockerVM.start_capture") as start_capture:
+
+            params = {"capture_file_name": "test.pcap", "data_link_type": "DLT_EN10MB"}
+            response = server.post("/projects/{project_id}/docker/vms/{vm_id}/adapters/0/ports/0/start_capture".format(project_id=vm["project_id"], vm_id=vm["vm_id"]), body=params, example=True)
+
+            assert response.status == 200
+
+            assert start_capture.called
+            assert "test.pcap" in response.json["pcap_file_path"]
+
+
+def test_docker_start_capture_not_started(server, vm, tmpdir):
+
+    with patch("gns3server.modules.docker.docker_vm.DockerVM.is_running", return_value=False) as mock:
+        with asyncio_patch("gns3server.modules.docker.docker_vm.DockerVM.start_capture") as start_capture:
+
+            params = {"capture_file_name": "test.pcap", "data_link_type": "DLT_EN10MB"}
+            response = server.post("/projects/{project_id}/docker/vms/{vm_id}/adapters/0/ports/0/start_capture".format(project_id=vm["project_id"], vm_id=vm["vm_id"]), body=params)
+
+            assert not start_capture.called
+            assert response.status == 409
+
+
+def test_docker_stop_capture(server, vm, tmpdir, project):
+
+    with patch("gns3server.modules.docker.docker_vm.DockerVM.is_running", return_value=True) as mock:
+        with asyncio_patch("gns3server.modules.docker.docker_vm.DockerVM.stop_capture") as stop_capture:
+
+            response = server.post("/projects/{project_id}/docker/vms/{vm_id}/adapters/0/ports/0/stop_capture".format(project_id=vm["project_id"], vm_id=vm["vm_id"]), example=True)
+
+            assert response.status == 204
+
+            assert stop_capture.called
+
+
+def test_docker_stop_capture_not_started(server, vm, tmpdir):
+
+    with patch("gns3server.modules.docker.docker_vm.DockerVM.is_running", return_value=False) as mock:
+        with asyncio_patch("gns3server.modules.docker.docker_vm.DockerVM.stop_capture") as stop_capture:
+
+            response = server.post("/projects/{project_id}/docker/vms/{vm_id}/adapters/0/ports/0/stop_capture".format(project_id=vm["project_id"], vm_id=vm["vm_id"]))
+
+            assert not stop_capture.called
+            assert response.status == 409

tests/modules/docker/test_docker_vm.py

@@ -133,7 +133,7 @@ def test_create_environment(loop, project, manager):
                 "Env": [
                     "YES=1",
                     "NO=0"
-                ],
+                    ],
                 "NetworkDisabled": True,
                 "Name": "test",
                 "Image": "ubuntu"
@@ -259,9 +259,6 @@ def test_start_without_nio(loop, vm, manager, free_console_port):
     assert vm.status != "started"
     vm.adapters = 1
 
-    # nio = manager.create_nio(0, {"type": "nio_udp", "lport": free_console_port, "rport": free_console_port, "rhost": "127.0.0.1"})
-    # loop.run_until_complete(asyncio.async(vm.adapter_add_nio_binding(0, nio)))
-
     with asyncio_patch("gns3server.modules.docker.DockerVM._get_container_state", return_value="stopped"):
         with asyncio_patch("gns3server.modules.docker.Docker.query") as mock_query:
             with asyncio_patch("gns3server.modules.docker.DockerVM._start_ubridge") as mock_start_ubridge:
@@ -544,3 +541,23 @@ def test_pull_image(loop, vm):
     with asyncio_patch("gns3server.modules.docker.Docker.http_query", return_value=mock_query) as mock:
         images = loop.run_until_complete(asyncio.async(vm.pull_image("ubuntu")))
         mock.assert_called_with("POST", "images/create", params={"fromImage": "ubuntu"})
+
+
+def test_start_capture(vm, tmpdir, manager, free_console_port, loop):
+
+    output_file = str(tmpdir / "test.pcap")
+    nio = manager.create_nio(0, {"type": "nio_udp", "lport": free_console_port, "rport": free_console_port, "rhost": "127.0.0.1"})
+    loop.run_until_complete(asyncio.async(vm.adapter_add_nio_binding(0, nio)))
+    loop.run_until_complete(asyncio.async(vm.start_capture(0, output_file)))
+    assert vm._ethernet_adapters[0].get_nio(0).capturing
+
+
+def test_stop_capture(vm, tmpdir, manager, free_console_port, loop):
+
+    output_file = str(tmpdir / "test.pcap")
+    nio = manager.create_nio(0, {"type": "nio_udp", "lport": free_console_port, "rport": free_console_port, "rhost": "127.0.0.1"})
+    loop.run_until_complete(asyncio.async(vm.adapter_add_nio_binding(0, nio)))
+    loop.run_until_complete(vm.start_capture(0, output_file))
+    assert vm._ethernet_adapters[0].get_nio(0).capturing
+    loop.run_until_complete(asyncio.async(vm.stop_capture(0)))
+    assert vm._ethernet_adapters[0].get_nio(0).capturing is False

tests/modules/vmware/test_vmware_vm.py

@@ -49,3 +49,25 @@ def test_json(vm, tmpdir, project):
     project._path = str(tmpdir)
     vm._linked_clone = True
     assert vm.__json__()["vm_directory"] is not None
+
+
+def test_start_capture(vm, tmpdir, manager, free_console_port, loop):
+
+    output_file = str(tmpdir / "test.pcap")
+    nio = manager.create_nio(0, {"type": "nio_udp", "lport": free_console_port, "rport": free_console_port, "rhost": "127.0.0.1"})
+    vm.adapters = 1
+    loop.run_until_complete(asyncio.async(vm.adapter_add_nio_binding(0, nio)))
+    loop.run_until_complete(asyncio.async(vm.start_capture(0, output_file)))
+    assert vm._ethernet_adapters[0].get_nio(0).capturing
+
+
+def test_stop_capture(vm, tmpdir, manager, free_console_port, loop):
+
+    output_file = str(tmpdir / "test.pcap")
+    nio = manager.create_nio(0, {"type": "nio_udp", "lport": free_console_port, "rport": free_console_port, "rhost": "127.0.0.1"})
+    vm.adapters = 1
+    loop.run_until_complete(asyncio.async(vm.adapter_add_nio_binding(0, nio)))
+    loop.run_until_complete(vm.start_capture(0, output_file))
+    assert vm._ethernet_adapters[0].get_nio(0).capturing
+    loop.run_until_complete(asyncio.async(vm.stop_capture(0)))
+    assert vm._ethernet_adapters[0].get_nio(0).capturing is False