Method for start / stop capture on a link

Ref https://github.com/GNS3/gns3-gui/issues/1117

gns3server/controller/link.py

@@ -15,16 +15,18 @@
 # 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 re
 import uuid
 import asyncio
 
 
 class Link:
 
-    def __init__(self, project):
+    def __init__(self, project, data_link_type="DLT_EN10MB"):
         self._id = str(uuid.uuid4())
         self._vms = []
         self._project = project
+        self._data_link_type = data_link_type
 
     @asyncio.coroutine
     def addVM(self, vm, adapter_number, port_number):
@@ -51,6 +53,35 @@ class Link:
         """
         raise NotImplementedError
 
+    @asyncio.coroutine
+    def start_capture(self):
+        """
+        Start capture on the link
+
+        :returns: Capture object
+        """
+        raise NotImplementedError
+
+    @asyncio.coroutine
+    def stop_capture(self):
+        """
+        Stop capture on the link
+        """
+        raise NotImplementedError
+
+    def capture_file_name(self):
+        """
+        :returns: File name for a capture on this link
+        """
+        capture_file_name = "{}_{}-{}_to_{}_{}-{}".format(
+            self._vms[0]["vm"].name,
+            self._vms[0]["adapter_number"],
+            self._vms[0]["port_number"],
+            self._vms[1]["vm"].name,
+            self._vms[1]["adapter_number"],
+            self._vms[1]["port_number"])
+        return re.sub("[^0-9A-Za-z_-]", "", capture_file_name) + ".pcap"
+
     @property
     def id(self):
         return self._id
@@ -63,4 +94,4 @@ class Link:
                 "adapter_number": side["adapter_number"],
                 "port_number": side["port_number"]
             })
-        return {"vms": res, "link_id": self._id}
+        return {"vms": res, "link_id": self._id, "data_link_type": self._data_link_type}

gns3server/controller/udp_link.py

@@ -16,6 +16,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import asyncio
+import aiohttp
 
 
 from .link import Link
@@ -23,8 +24,9 @@ from .link import Link
 
 class UDPLink(Link):
 
-    def __init__(self, project):
-        super().__init__(project)
+    def __init__(self, project, data_link_type="DLT_EN10MB"):
+        super().__init__(project, data_link_type)
+        self._capture_vm = None
 
     @asyncio.coroutine
     def create(self):
@@ -76,3 +78,45 @@ class UDPLink(Link):
 
         yield from vm1.delete("/adapters/{adapter_number}/ports/{port_number}/nio".format(adapter_number=adapter_number1, port_number=port_number1))
         yield from vm2.delete("/adapters/{adapter_number}/ports/{port_number}/nio".format(adapter_number=adapter_number2, port_number=port_number2))
+
+    @asyncio.coroutine
+    def start_capture(self):
+        """
+        Start capture on a link
+        """
+        self._capture_vm = self._choose_capture_side()
+        data = {
+            "capture_file_name": self.capture_file_name(),
+            "data_link_type": self._data_link_type
+        }
+        yield from self._capture_vm["vm"].post("/adapters/{adapter_number}/ports/{port_number}/start_capture".format(adapter_number=self._capture_vm["adapter_number"], port_number=self._capture_vm["port_number"]), data=data)
+
+    @asyncio.coroutine
+    def stop_capture(self):
+        """
+        Stop capture on a link
+        """
+        if self._capture_vm:
+            yield from self._capture_vm["vm"].post("/adapters/{adapter_number}/ports/{port_number}/stop_capture".format(adapter_number=self._capture_vm["adapter_number"], port_number=self._capture_vm["port_number"]))
+            self._capture_vm = None
+
+    def _choose_capture_side(self):
+        """
+        Run capture on the best candidate.
+
+        The ideal candidate is a node who support capture on controller
+        server
+
+        :returns: VM where the capture should run
+        """
+
+        # For saving bandwith we use the local node first
+        for vm in self._vms:
+            if vm["vm"].compute.id == "local" and vm["vm"].vm_type not in ["qemu", "vpcs"]:
+                return vm
+
+        for vm in self._vms:
+            if vm["vm"].vm_type not in ["qemu", "vpcs"]:
+                return vm
+
+        raise aiohttp.web.HTTPConflict(text="Capture is not supported for this link")

gns3server/controller/vm.py

@@ -216,6 +216,9 @@ class VM:
         else:
             return (yield from self._compute.delete("/projects/{}/{}/vms/{}{}".format(self._project.id, self._vm_type, self._id, path)))
 
+    def __repr__(self):
+        return "<gns3server.controller.VM {} {}>".format(self._vm_type, self._name)
+
     def __json__(self):
         return {
             "compute_id": self._compute.id,

gns3server/handlers/api/controller/link_handler.py

@@ -52,6 +52,46 @@ class LinkHandler:
         response.set_status(201)
         response.json(link)
 
+    @classmethod
+    @Route.post(
+        r"/projects/{project_id}/links/{link_id}/start_capture",
+        parameters={
+            "project_id": "UUID for the project",
+            "link_id": "UUID of the link"
+        },
+        status_codes={
+            204: "Capture started",
+            400: "Invalid request"
+        },
+        description="Start capture on a link instance")
+    def start_capture(request, response):
+
+        controller = Controller.instance()
+        project = controller.getProject(request.match_info["project_id"])
+        link = project.getLink(request.match_info["link_id"])
+        yield from link.start_capture()
+        response.set_status(204)
+
+    @classmethod
+    @Route.post(
+        r"/projects/{project_id}/links/{link_id}/stop_capture",
+        parameters={
+            "project_id": "UUID for the project",
+            "link_id": "UUID of the link"
+        },
+        status_codes={
+            204: "Capture stopped",
+            400: "Invalid request"
+        },
+        description="Stop capture on a link instance")
+    def stop_capture(request, response):
+
+        controller = Controller.instance()
+        project = controller.getProject(request.match_info["project_id"])
+        link = project.getLink(request.match_info["link_id"])
+        yield from link.stop_capture()
+        response.set_status(204)
+
     @classmethod
     @Route.delete(
         r"/projects/{project_id}/links/{link_id}",

gns3server/schemas/link.py

@@ -28,6 +28,10 @@ LINK_OBJECT_SCHEMA = {
             "maxLength": 36,
             "pattern": "^[a-fA-F0-9]{8}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{4}-[a-fA-F0-9]{12}$"
         },
+        "data_link_type": {
+            "description": "PCAP data link type (http://www.tcpdump.org/linktypes.html)",
+            "enum": ["DLT_ATM_RFC1483", "DLT_EN10MB", "DLT_FRELAY", "DLT_C_HDLC"]
+        },
         "vms": {
             "description": "List of the VMS",
             "type": "array",

gns3server/schemas/vm.py

@@ -54,9 +54,8 @@ VM_CAPTURE_SCHEMA = {
             "minLength": 1,
         },
         "data_link_type": {
-            "description": "PCAP data link type",
-            "type": "string",
-            "minLength": 1,
+            "description": "PCAP data link type (http://www.tcpdump.org/linktypes.html)",
+            "enum": ["DLT_ATM_RFC1483", "DLT_EN10MB", "DLT_FRELAY", "DLT_C_HDLC"]
         }
     },
     "additionalProperties": False,

tests/controller/test_link.py

@@ -57,6 +57,7 @@ def test_json(async_run, project, compute):
     async_run(link.addVM(vm2, 1, 3))
     assert link.__json__() == {
         "link_id": link.id,
+        "data_link_type": "DLT_EN10MB",
         "vms": [
             {
                 "vm_id": vm1.id,
@@ -70,3 +71,12 @@ def test_json(async_run, project, compute):
             }
         ]
     }
+
+def test_capture_filename(project, compute, async_run):
+    vm1 = VM(project, compute, name="Hello@")
+    vm2 = VM(project, compute, name="w0.rld")
+
+    link = Link(project)
+    async_run(link.addVM(vm1, 0, 4))
+    async_run(link.addVM(vm2, 1, 3))
+    assert link.capture_file_name() == "Hello_0-4_to_w0rld_1-3.pcap"

tests/controller/test_udp_link.py

@@ -97,3 +97,65 @@ def test_delete(async_run, project):
 
     compute1.delete.assert_any_call("/projects/{}/vpcs/vms/{}/adapters/0/ports/4/nio".format(project.id, vm1.id))
     compute2.delete.assert_any_call("/projects/{}/vpcs/vms/{}/adapters/3/ports/1/nio".format(project.id, vm2.id))
+
+
+def test_choose_capture_side(async_run, project):
+    """
+    The link capture should run on the optimal node
+    """
+    compute1 = MagicMock()
+    compute2 = MagicMock()
+    compute2.id = "local"
+
+    vm_vpcs = VM(project, compute1, vm_type="vpcs")
+    vm_iou = VM(project, compute2, vm_type="iou")
+
+    link = UDPLink(project)
+    async_run(link.addVM(vm_vpcs, 0, 4))
+    async_run(link.addVM(vm_iou, 3, 1))
+
+    assert link._choose_capture_side()["vm"] == vm_iou
+
+    vm_vpcs = VM(project, compute1, vm_type="vpcs")
+    vm_vpcs2 = VM(project, compute1, vm_type="vpcs")
+
+    link = UDPLink(project)
+    async_run(link.addVM(vm_vpcs, 0, 4))
+    async_run(link.addVM(vm_vpcs2, 3, 1))
+
+    # VPCS doesn't support capture
+    with pytest.raises(aiohttp.web.HTTPConflict):
+        link._choose_capture_side()["vm"]
+
+    # Capture should run on the local node
+    vm_iou = VM(project, compute1, vm_type="iou")
+    vm_iou2 = VM(project, compute2, vm_type="iou")
+
+    link = UDPLink(project)
+    async_run(link.addVM(vm_iou, 0, 4))
+    async_run(link.addVM(vm_iou2, 3, 1))
+
+    assert link._choose_capture_side()["vm"] == vm_iou2
+
+
+def test_capture(async_run, project):
+    compute1 = MagicMock()
+
+    vm_vpcs = VM(project, compute1, vm_type="vpcs", name="V1")
+    vm_iou = VM(project, compute1, vm_type="iou", name="I1")
+
+    link = UDPLink(project)
+    async_run(link.addVM(vm_vpcs, 0, 4))
+    async_run(link.addVM(vm_iou, 3, 1))
+
+    capture = async_run(link.start_capture())
+
+    compute1.post.assert_any_call("/projects/{}/iou/vms/{}/adapters/3/ports/1/start_capture".format(project.id, vm_iou.id), data={
+        "capture_file_name": link.capture_file_name(),
+        "data_link_type": "DLT_EN10MB"
+    })
+
+    capture = async_run(link.stop_capture())
+
+    compute1.post.assert_any_call("/projects/{}/iou/vms/{}/adapters/3/ports/1/stop_capture".format(project.id, vm_iou.id))
+

tests/handlers/api/controller/test_link.py

@@ -36,11 +36,11 @@ from gns3server.controller.link import Link
 
 
 @pytest.fixture
-def hypervisor(http_controller, async_run):
-    hypervisor = MagicMock()
-    hypervisor.id = "example.com"
-    Controller.instance()._hypervisors = {"example.com": hypervisor}
-    return hypervisor
+def compute(http_controller, async_run):
+    compute = MagicMock()
+    compute.id = "example.com"
+    Controller.instance()._computes = {"example.com": compute}
+    return compute
 
 
 @pytest.fixture
@@ -48,13 +48,13 @@ def project(http_controller, async_run):
     return async_run(Controller.instance().addProject())
 
 
-def test_create_link(http_controller, tmpdir, project, hypervisor, async_run):
+def test_create_link(http_controller, tmpdir, project, compute, async_run):
     response = MagicMock()
     response.json = {"console": 2048}
-    hypervisor.post = AsyncioMagicMock(return_value=response)
+    compute.post = AsyncioMagicMock(return_value=response)
 
-    vm1 = async_run(project.addVM(hypervisor, None))
-    vm2 = async_run(project.addVM(hypervisor, None))
+    vm1 = async_run(project.addVM(compute, None))
+    vm2 = async_run(project.addVM(compute, None))
 
     with asyncio_patch("gns3server.controller.udp_link.UDPLink.create") as mock:
         response = http_controller.post("/projects/{}/links".format(project.id), {
@@ -77,10 +77,29 @@ def test_create_link(http_controller, tmpdir, project, hypervisor, async_run):
     assert len(response.json["vms"]) == 2
 
 
-def test_delete_link(http_controller, tmpdir, project, hypervisor, async_run):
+def test_start_capture(http_controller, tmpdir, project, compute, async_run):
+    link = Link(project)
+    project._links = {link.id: link}
+    with asyncio_patch("gns3server.controller.link.Link.start_capture") as mock:
+        response = http_controller.post("/projects/{}/links/{}/start_capture".format(project.id, link.id), example=True)
+    assert mock.called
+    assert response.status == 204
 
+
+def test_stop_capture(http_controller, tmpdir, project, compute, async_run):
     link = Link(project)
     project._links = {link.id: link}
-    with asyncio_patch("gns3server.controller.udp_link.Link.delete"):
+    with asyncio_patch("gns3server.controller.link.Link.stop_capture") as mock:
+        response = http_controller.post("/projects/{}/links/{}/stop_capture".format(project.id, link.id), example=True)
+    assert mock.called
+    assert response.status == 204
+
+
+def test_delete_link(http_controller, tmpdir, project, compute, async_run):
+
+    link = Link(project)
+    project._links = {link.id: link}
+    with asyncio_patch("gns3server.controller.link.Link.delete") as mock:
         response = http_controller.delete("/projects/{}/links/{}".format(project.id, link.id), example=True)
+    assert mock.called
     assert response.status == 204