Fix download on 3.x (presigned URL)

Author: andreax79

airflow_code_editor/api/api.py

@@ -26,11 +26,13 @@
 from airflow_code_editor.commons import (
     HTTP_200_OK,
     HTTP_400_BAD_REQUEST,
+    HTTP_401_UNAUTHORIZED,
     HTTP_404_NOT_FOUND,
     HTTP_500_SERVER_ERROR,
     VERSION,
 )
 from airflow_code_editor.fs import RootFS
+from airflow_code_editor.presigned import create_presigned, decode_presigned
 from airflow_code_editor.tree import get_stat, get_tree
 from airflow_code_editor.utils import (
     DummyLexer,
@@ -53,6 +55,8 @@
     "search",
     "ping",
     "get_version",
+    "generate_presigned",
+    "load_presigned",
 ]
 
 
@@ -250,3 +254,23 @@ def get_version():
         "version": VERSION,
         "airflow_version": airflow_version,
     }
+
+
+def generate_presigned(path):
+    "Generate a presigned URL token for downloading a file/git object"
+    return {
+        "token": create_presigned(path),
+    }
+
+
+def load_presigned(token: str):
+    "Download a file/git object using a presigned URL"
+    try:
+        path = decode_presigned(token)
+    except Exception as ex:
+        logging.error(ex)
+        return prepare_api_response(
+            error_message="Not authenticated",
+            http_status_code=HTTP_401_UNAUTHORIZED,
+        )
+    return load(path)

airflow_code_editor/api/code_editor.yaml

@@ -1,7 +1,7 @@
 openapi: "3.0.3"
 info:
   title: "Airflow Code Editor API"
-  version: "1.0.0"
+  version: "8.1.0"
   description: |
     # Overview
 
@@ -218,6 +218,56 @@ paths:
           description: "Success"
           $ref: '#/components/responses/VersionInfo'
 
+  /generate_presigned:
+    post:
+      summary: "Generate a presigned URL token for downloading a file/git object"
+      x-openapi-router-controller: airflow_code_editor.api.flask_endpoints
+      operationId: generate_presigned
+      tags: [Presigned]
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              type: object
+              properties:
+                path:
+                  type: string
+                  description: "File/git object path"
+      responses:
+        '200':
+          description: "Success"
+          $ref: '#/components/responses/GeneratePresignedResponse'
+        '401':
+          description: "Not authenticated"
+        '403':
+          description: "Client does not have sufficient permission"
+
+  /presigned/{token}:
+    get:
+      summary: "Download a file/git object using a presigned URL"
+      x-openapi-router-controller: airflow_code_editor.api.flask_endpoints
+      operationId: load_presigned
+      tags: [Presigned]
+      parameters:
+        - name: token
+          in: path
+          description: "Presigned URL token"
+          required: true
+          allowReserved: true
+          schema:
+            type: string
+            format: path
+      responses:
+        '200':
+          description: "Success"
+        '403':
+          description: "Client does not have sufficient permission"
+        '401':
+          description: "Not authenticated"
+        '404':
+          description: "File not found"
+
 components:
   schemas:
     TreeEntity:
@@ -328,6 +378,19 @@ components:
               - version
               - airflow_version
 
+    GeneratePresignedResponse:
+      description: Generate presigned URL token response
+      content:
+        application/json:
+          schema:
+            type: object
+            properties:
+              token:
+                type: string
+                description: Presigned URL token
+            required:
+              - token
+
   parameters:
     All:
       in: query

airflow_code_editor/api/fastapi_endpoints.py

@@ -71,6 +71,28 @@ def delete(path: str, request: Request):
     return api.delete(path)
 
 
+@app.post(
+    "/generate_presigned",
+    dependencies=[Depends(requires_access_dag(method="GET"))],
+    include_in_schema=False,
+)
+async def generate_presigned(request: Request):
+    "Generate a presigned URL token for downloading a file/git object"
+    body = await request.json()
+    path = body.get("path", "")
+    return api.generate_presigned(path)
+
+
+@app.get(
+    "/presigned/{token}",
+    dependencies=[],  # presigned URL does not require authentication
+    include_in_schema=False,
+)
+def presigned(token: str, request: Request):
+    "Download a file/git object using a presigned URL"
+    return api.load_presigned(token)
+
+
 @app.post(
     "/format",
     dependencies=[Depends(requires_access_dag(method="GET"))],

airflow_code_editor/api/flask_endpoints.py

@@ -35,6 +35,8 @@
     "search",
     "post_git",
     "get_version",
+    "generate_presigned",
+    "load_presigned",
     "load_specification",
     "api_blueprint",
 ]
@@ -99,6 +101,21 @@ def get_version():
     return api.get_version()
 
 
+@security.requires_access_dag("PUT")
+@csrf.exempt
+def generate_presigned(*, path: str = None):
+    "Generate a presigned URL token for downloading a file/git object"
+    path = request.json.get("path", "")
+    return api.generate_presigned(path)
+
+
+# security is not required for presigned URLs
+@csrf.exempt
+def load_presigned(*, token: str = None):
+    "Download a file/git object using a presigned URL"
+    return api.load_presigned(token)
+
+
 def load_specification() -> dict:
     with importlib.resources.path("airflow_code_editor.api", "code_editor.yaml") as f:
         return safe_load(f.read_text())

airflow_code_editor/app_builder_view.py

@@ -124,6 +124,18 @@ def get_version(self):
     def ping(self):
         return api.ping()
 
+    @expose("/generate_presigned", methods=["POST"])
+    @auth.has_access(PERMISSIONS)
+    def generate_presigned(self):
+        path = request.json.get("path", "")
+        return api.generate_presigned(path)
+
+    @expose("/presigned/<path:path>", methods=["GET"])
+    # auth is not required for presigned URLs
+    def load_presigned(self, path=None):
+        "Download a file/git object using a presigned URL"
+        return api.load_presigned(path)
+
 
 appbuilder_code_editor_view = AppBuilderCodeEditorView()
 appbuilder_view = {

airflow_code_editor/commons.py

@@ -32,6 +32,7 @@
     'SUPPORTED_GIT_COMMANDS',
     'HTTP_200_OK',
     'HTTP_400_BAD_REQUEST',
+    'HTTP_401_UNAUTHORIZED',
     'HTTP_404_NOT_FOUND',
     'HTTP_500_SERVER_ERROR',
     'PLUGIN_DEFAULT_CONFIG',
@@ -64,6 +65,7 @@
 DEFAULT_GIT_BRANCH = 'main'
 HTTP_200_OK = 200
 HTTP_400_BAD_REQUEST = 400
+HTTP_401_UNAUTHORIZED = 401
 HTTP_404_NOT_FOUND = 404
 HTTP_500_SERVER_ERROR = 500
 SUPPORTED_GIT_COMMANDS = [

airflow_code_editor/fastapi_app.py

@@ -24,7 +24,7 @@
 from fastapi.templating import Jinja2Templates
 
 from airflow_code_editor.api.fastapi_endpoints import app
-from airflow_code_editor.commons import PLUGIN_LONG_NAME, PLUGIN_NAME, MENU_LABEL
+from airflow_code_editor.commons import MENU_LABEL, PLUGIN_LONG_NAME, PLUGIN_NAME
 from airflow_code_editor.utils import is_enabled
 
 __all__ = ["fastapi_app"]

airflow_code_editor/presigned.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+#
+#   Copyright 2019 Andrea Bonomi <[email protected]>
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License
+#
+
+import time
+
+from airflow.configuration import conf
+from itsdangerous import URLSafeSerializer
+
+__all__ = ["create_presigned", "decode_presigned"]
+
+
+def get_signer() -> URLSafeSerializer:
+    """
+    Get a signer instance
+    """
+    secret_key = conf.get_mandatory_value("core", "fernet_key")
+    return URLSafeSerializer(secret_key)
+
+
+def create_presigned(filename: str, expires_in: int = 300) -> str:
+    """
+    Generate a signed token with expiry
+    """
+    payload = {"filename": filename, "exp": int(time.time()) + expires_in}
+    return get_signer().dumps(payload)
+
+
+def decode_presigned(token: str) -> str:
+    """
+    Decode and verify a signed token
+    """
+    data = get_signer().loads(token)
+    if data["exp"] < int(time.time()):
+        raise ValueError("Token has expired")
+    return data["filename"]