fix(platform): Fix marketplace leaking secrets (#8281)

* add hide secrets param * Update autogpt_platform/backend/backend/data/graph.py Co-authored-by: Reinier van der Leer <[email protected]> * Update autogpt_platform/frontend/src/lib/autogpt-server-api/baseClient.ts Co-authored-by: Reinier van der Leer <[email protected]> * rename hide secrets * use builtin dict * delete key --------- Co-authored-by: Reinier van der Leer <[email protected]>
Significant-Gravitas · Oct 8, 2024 · bc1df92 · bc1df92
1 parent 04473ca
commit bc1df92
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 8 deletions.
diff --git a/autogpt_platform/backend/backend/data/graph.py b/autogpt_platform/backend/backend/data/graph.py
@@ -330,7 +330,7 @@ def get_input_schema(self) -> list[InputSchemaItem]:
         return input_schema
 
     @staticmethod
-    def from_db(graph: AgentGraph):
+    def from_db(graph: AgentGraph, hide_credentials: bool = False):
         nodes = [
             *(graph.AgentNodes or []),
             *(
@@ -341,7 +341,7 @@ def from_db(graph: AgentGraph):
         ]
         return Graph(
             **GraphMeta.from_db(graph).model_dump(),
-            nodes=[Node.from_db(node) for node in nodes],
+            nodes=[Graph._process_node(node, hide_credentials) for node in nodes],
             links=list(
                 {
                     Link.from_db(link)
@@ -355,6 +355,26 @@ def from_db(graph: AgentGraph):
             },
         )
 
+    @staticmethod
+    def _process_node(node: AgentNode, hide_credentials: bool) -> Node:
+        node_dict = node.model_dump()
+        if hide_credentials and "constantInput" in node_dict:
+            constant_input = json.loads(node_dict["constantInput"])
+            Graph._hide_credentials_in_input(constant_input)
+            node_dict["constantInput"] = json.dumps(constant_input)
+        return Node.from_db(AgentNode(**node_dict))
+
+    @staticmethod
+    def _hide_credentials_in_input(input_data: dict[str, Any]):
+        sensitive_keys = ["credentials", "api_key", "password", "token", "secret"]
+        for key, value in input_data.items():
+            if isinstance(value, dict):
+                Graph._hide_credentials_in_input(value)
+            elif isinstance(value, str) and any(
+                sensitive_key in key.lower() for sensitive_key in sensitive_keys
+            ):
+                del input_data[key]
+
 
 AGENT_NODE_INCLUDE: prisma.types.AgentNodeInclude = {
     "Input": True,
@@ -431,6 +451,7 @@ async def get_graph(
     version: int | None = None,
     template: bool = False,
     user_id: str | None = None,
+    hide_credentials: bool = False,
 ) -> Graph | None:
     """
     Retrieves a graph from the DB.
@@ -456,7 +477,7 @@ async def get_graph(
         include=AGENT_GRAPH_INCLUDE,
         order={"version": "desc"},
     )
-    return Graph.from_db(graph) if graph else None
+    return Graph.from_db(graph, hide_credentials) if graph else None
 
 
 async def set_graph_active_version(graph_id: str, version: int, user_id: str) -> None:

diff --git a/autogpt_platform/backend/backend/server/rest_api.py b/autogpt_platform/backend/backend/server/rest_api.py
@@ -364,8 +364,11 @@ async def get_graph(
         graph_id: str,
         user_id: Annotated[str, Depends(get_user_id)],
         version: int | None = None,
+        hide_credentials: bool = False,
     ) -> graph_db.Graph:
-        graph = await graph_db.get_graph(graph_id, version, user_id=user_id)
+        graph = await graph_db.get_graph(
+            graph_id, version, user_id=user_id, hide_credentials=hide_credentials
+        )
         if not graph:
             raise HTTPException(status_code=404, detail=f"Graph #{graph_id} not found.")
         return graph

diff --git a/autogpt_platform/frontend/src/app/marketplace/submit/page.tsx b/autogpt_platform/frontend/src/app/marketplace/submit/page.tsx
@@ -129,7 +129,7 @@ const SubmitPage: React.FC = () => {
     const fetchAgentGraph = async () => {
       if (selectedAgentId) {
         const api = new AutoGPTServerAPI();
-        const graph = await api.getGraph(selectedAgentId);
+        const graph = await api.getGraph(selectedAgentId, undefined, true);
         setSelectedAgentGraph(graph);
         setValue("name", graph.name);
         setValue("description", graph.description);

diff --git a/autogpt_platform/frontend/src/lib/autogpt-server-api/baseClient.ts b/autogpt_platform/frontend/src/lib/autogpt-server-api/baseClient.ts
@@ -70,9 +70,19 @@ export default class BaseAutoGPTServerAPI {
     return this._get("/templates");
   }
 
-  getGraph(id: string, version?: number): Promise<Graph> {
-    const query = version !== undefined ? `?version=${version}` : "";
-    return this._get(`/graphs/${id}` + query);
+  getGraph(
+    id: string,
+    version?: number,
+    hide_credentials?: boolean,
+  ): Promise<Graph> {
+    let query: Record<string, any> = {};
+    if (version !== undefined) {
+      query["version"] = version;
+    }
+    if (hide_credentials !== undefined) {
+      query["hide_credentials"] = hide_credentials;
+    }
+    return this._get(`/graphs/${id}`, query);
   }
 
   getTemplate(id: string, version?: number): Promise<Graph> {