From 46c10c030eb1e53e2bdc420278681ebd1b44d4a9 Mon Sep 17 00:00:00 2001
From: Matteo Luppi <100372313+lups2000@users.noreply.github.com>
Date: Sun, 1 Sep 2024 22:22:51 +0200
Subject: [PATCH] add 2 new endpoints for processes extraction and process
 image (#7136)

* add 2 new endpoints for processes extraction and process image

* [autofix.ci] apply automated fixes

* add review changes

* add tests

* [autofix.ci] apply automated fixes

* nit

* update tests

* [autofix.ci] apply automated fixes

* add fallback image and update tests

* [autofix.ci] apply automated fixes

* fix lint error

* fix tests

* [autofix.ci] apply automated fixes

* try to use base64

* still trying to fix test on win

* [autofix.ci] apply automated fixes

* nit

* [autofix.ci] apply automated fixes

* TRANSPARENT_PNG: use raw bytes to avoid base64 step

* tests: use feature-based detection

* hardening: prevent mime type sniffing

* fixup feature detection

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Maximilian Hils <git@maximilianhils.com>
---
 mitmproxy/tools/web/app.py           | 44 ++++++++++++++++++++++++++++
 pyproject.toml                       |  2 +-
 test/mitmproxy/tools/web/test_app.py | 28 ++++++++++++++++++
 3 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/mitmproxy/tools/web/app.py b/mitmproxy/tools/web/app.py
index 59fa5674c..9cb42c620 100644
--- a/mitmproxy/tools/web/app.py
+++ b/mitmproxy/tools/web/app.py
@@ -13,6 +13,7 @@ from io import BytesIO
 from itertools import islice
 from typing import ClassVar
 
+import mitmproxy_rs
 import tornado.escape
 import tornado.web
 import tornado.websocket
@@ -38,6 +39,12 @@ from mitmproxy.utils.emoji import emoji
 from mitmproxy.utils.strutils import always_str
 from mitmproxy.websocket import WebSocketMessage
 
+TRANSPARENT_PNG = (
+    b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08"
+    b"\x04\x00\x00\x00\xb5\x1c\x0c\x02\x00\x00\x00\x0bIDATx\xdac\xfc\xff\x07"
+    b"\x00\x02\x00\x01\xfc\xa8Q\rh\x00\x00\x00\x00IEND\xaeB`\x82"
+)
+
 
 def cert_to_json(certs: Sequence[certs.Cert]) -> dict | None:
     if not certs:
@@ -654,6 +661,41 @@ class State(RequestHandler):
         self.write(State.get_json(self.master))
 
 
+class ProcessList(RequestHandler):
+    @staticmethod
+    def get_json():
+        processes = mitmproxy_rs.active_executables()
+        return [
+            {
+                "is_visible": process.is_visible,
+                "executable": process.executable,
+                "is_system": process.is_system,
+                "display_name": process.display_name,
+            }
+            for process in processes
+        ]
+
+    def get(self):
+        self.write(ProcessList.get_json())
+
+
+class ProcessImage(RequestHandler):
+    def get(self):
+        path = self.get_query_argument("path", None)
+
+        if not path:
+            raise APIError(400, "Missing 'path' parameter.")
+
+        try:
+            icon_bytes = mitmproxy_rs.executable_icon(path)
+        except Exception:
+            icon_bytes = TRANSPARENT_PNG
+
+        self.set_header("Content-Type", "image/png")
+        self.set_header("X-Content-Type-Options", "nosniff")
+        self.write(icon_bytes)
+
+
 class GZipContentAndFlowFiles(tornado.web.GZipContentEncoding):
     CONTENT_TYPES = {
         "application/octet-stream",
@@ -713,5 +755,7 @@ class Application(tornado.web.Application):
                 (r"/options(?:\.json)?", Options),
                 (r"/options/save", SaveOptions),
                 (r"/state(?:\.json)?", State),
+                (r"/processes", ProcessList),
+                (r"/executable-icon", ProcessImage),
             ],
         )
diff --git a/pyproject.toml b/pyproject.toml
index 703c57c4c..0ac709a66 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -43,7 +43,7 @@ dependencies = [
     "hyperframe>=6.0,<=6.0.1",
     "kaitaistruct>=0.10,<=0.10",
     "ldap3>=2.8,<=2.9.1",
-    "mitmproxy_rs>=0.7.1,<0.8",  # relaxed upper bound here: we control this
+    "mitmproxy_rs>=0.7.2,<0.8",  # relaxed upper bound here: we control this
     "msgpack>=1.0.0,<=1.0.8",
     "passlib>=1.6.5,<=1.7.4",
     "protobuf>=5.27.2,<=5.27.3",
diff --git a/test/mitmproxy/tools/web/test_app.py b/test/mitmproxy/tools/web/test_app.py
index 2e9b7f36e..f5677e404 100644
--- a/test/mitmproxy/tools/web/test_app.py
+++ b/test/mitmproxy/tools/web/test_app.py
@@ -5,6 +5,7 @@ import logging
 from pathlib import Path
 from unittest import mock
 
+import mitmproxy_rs
 import pytest
 import tornado.testing
 from tornado import httpclient
@@ -403,3 +404,30 @@ class TestApp(tornado.testing.AsyncHTTPTestCase):
         # trigger on_close by opening a second connection.
         ws_client2 = yield websocket.websocket_connect(ws_url)
         ws_client2.close()
+
+    def test_process_list(self):
+        try:
+            mitmproxy_rs.active_executables()
+        except NotImplementedError:
+            pytest.skip(
+                "mitmproxy_rs.active_executables not available on this platform."
+            )
+        resp = self.fetch("/processes")
+        assert resp.code == 200
+        assert get_json(resp)
+
+    def test_process_icon(self):
+        try:
+            mitmproxy_rs.executable_icon("invalid")
+        except NotImplementedError:
+            pytest.skip("mitmproxy_rs.executable_icon not available on this platform.")
+        except Exception:
+            pass
+        resp = self.fetch("/executable-icon")
+        assert resp.code == 400
+        assert "Missing 'path' parameter." in resp.body.decode()
+
+        resp = self.fetch("/executable-icon?path=invalid_path")
+        assert resp.code == 200
+        assert resp.headers["Content-Type"] == "image/png"
+        assert resp.body == app.TRANSPARENT_PNG