From be6d4771d097121e9f108a196bdea779a165a627 Mon Sep 17 00:00:00 2001
From: grossmj <grossmj@gns3.net>
Date: Mon, 19 Oct 2020 15:55:29 +1030
Subject: [PATCH] Migrate PCAP streaming code to work with FastAPI.

---
 gns3server/endpoints/controller/links.py | 67 +++++++++++-------------
 gns3server/endpoints/index.py            |  6 +--
 2 files changed, 35 insertions(+), 38 deletions(-)

diff --git a/gns3server/endpoints/controller/links.py b/gns3server/endpoints/controller/links.py
index 3b52b8b9..209c261a 100644
--- a/gns3server/endpoints/controller/links.py
+++ b/gns3server/endpoints/controller/links.py
@@ -19,9 +19,12 @@
 API endpoints for links.
 """
 
+import aiohttp
+import multidict
+
 from fastapi import APIRouter, Depends, Request, status
-from fastapi.encoders import jsonable_encoder
 from fastapi.responses import StreamingResponse
+from fastapi.encoders import jsonable_encoder
 from typing import List
 from uuid import UUID
 
@@ -181,37 +184,31 @@ async def reset_link(link: Link = Depends(dep_link)):
     return link.__json__()
 
 
-# @router.post("/projects/{project_id}/links/{link_id}/pcap",
-#              summary="Stream a packet capture",
-#              responses={404: {"model": ErrorMessage, "description": "Project or link not found"}})
-# async def pcap(project_id: UUID, link_id: UUID, request: Request):
-#     """
-#     Stream the PCAP capture file from compute.
-#     """
-#
-#     project = await Controller.instance().get_loaded_project(str(project_id))
-#     link = project.get_link(str(link_id))
-#     if not link.capturing:
-#         raise ControllerError("This link has no active packet capture")
-#
-#     compute = link.compute
-#     pcap_streaming_url = link.pcap_streaming_url()
-#     headers = multidict.MultiDict(request.headers)
-#     headers['Host'] = compute.host
-#     headers['Router-Host'] = request.client.host
-#     body = await request.body()
-#
-#     connector = aiohttp.TCPConnector(limit=None, force_close=True)
-#     async with aiohttp.ClientSession(connector=connector, headers=headers) as session:
-#         async with session.request(request.method, pcap_streaming_url, timeout=None, data=body) as response:
-#             proxied_response = aiohttp.web.Response(headers=response.headers, status=response.status)
-#             if response.headers.get('Transfer-Encoding', '').lower() == 'chunked':
-#                 proxied_response.enable_chunked_encoding()
-#
-#             await proxied_response.prepare(request)
-#             async for data in response.content.iter_any():
-#                 if not data:
-#                     break
-#                 await proxied_response.write(data)
-#
-#     #return StreamingResponse(file_like, media_type="video/mp4"))
+@router.get("/{link_id}/pcap",
+            responses=responses)
+async def pcap(request: Request, link: Link = Depends(dep_link)):
+    """
+    Stream the PCAP capture file from compute.
+    """
+
+    if not link.capturing:
+        raise ControllerError("This link has no active packet capture")
+
+    compute = link.compute
+    pcap_streaming_url = link.pcap_streaming_url()
+    headers = multidict.MultiDict(request.headers)
+    headers['Host'] = compute.host
+    headers['Router-Host'] = request.client.host
+    body = await request.body()
+
+    async def compute_pcpa_stream():
+
+        connector = aiohttp.TCPConnector(limit=None, force_close=True)
+        async with aiohttp.ClientSession(connector=connector, headers=headers) as session:
+            async with session.request(request.method, pcap_streaming_url, timeout=None, data=body) as compute_response:
+                async for data in compute_response.content.iter_any():
+                    if not data:
+                        break
+                    yield data
+
+    return StreamingResponse(compute_pcpa_stream(), media_type="application/vnd.tcpdump.pcap")
diff --git a/gns3server/endpoints/index.py b/gns3server/endpoints/index.py
index 21e83179..df932017 100644
--- a/gns3server/endpoints/index.py
+++ b/gns3server/endpoints/index.py
@@ -16,7 +16,7 @@
 
 import os
 
-from fastapi import APIRouter, Request, HTTPException
+from fastapi import APIRouter, Request, HTTPException, status
 from fastapi.responses import RedirectResponse, HTMLResponse, FileResponse
 from fastapi.templating import Jinja2Templates
 
@@ -30,7 +30,7 @@ templates = Jinja2Templates(directory=os.path.join("gns3server", "templates"))
 @router.get("/")
 async def root():
 
-    return RedirectResponse("/static/web-ui/bundled", status_code=308)
+    return RedirectResponse("/static/web-ui/bundled", status_code=308)  # permanent redirect
 
 
 @router.get("/debug",
@@ -54,7 +54,7 @@ async def web_ui(file_path: str):
 
     # Raise error if user try to escape
     if file_path[0] == ".":
-        raise HTTPException(status_code=403)
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN)
 
     static = get_resource(file_path)