From 1b105b4746d6ff893a6d9cc388d88149d15b106a Mon Sep 17 00:00:00 2001
From: Jeroen Beckers <info@dauntless.be>
Date: Sat, 27 Jul 2024 10:30:01 +0100
Subject: [PATCH] Add Universal edit button (#2660)

* Clean base

* Remove edit btn from a few pages

* Update path logic

* Add hook to mkdocs

* Fix universal edit button functionality and restrict for specific file paths

* Refactor edit button URI handling code into a function

---------

Co-authored-by: Jeroen Beckers <dauntless@dauntless.be>
Co-authored-by: Carlos Holguera <perezholguera@gmail.com>
---
 docs/hooks/edit_button_uri.py | 43 +++++++++++++++++++++++++++++++++++
 mkdocs.yml                    |  5 +++-
 2 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 docs/hooks/edit_button_uri.py

diff --git a/docs/hooks/edit_button_uri.py b/docs/hooks/edit_button_uri.py
new file mode 100644
index 0000000000..fa2623bd92
--- /dev/null
+++ b/docs/hooks/edit_button_uri.py
@@ -0,0 +1,43 @@
+import logging
+
+log = logging.getLogger('mkdocs')
+
+def get_edit_url(src_path, edit_url_mastg, edit_url_masvs):
+    if src_path.startswith("MASVS"):
+        edit_url = f"{edit_url_masvs}{src_path}"
+        edit_url = edit_url.replace("master/MASVS/controls", "master/controls/")
+        edit_url = edit_url.replace("master/MASVS/", "master/Document/")
+    elif src_path.startswith("MASTG"):
+        edit_url = f"{edit_url_mastg}{src_path}"
+        edit_url = edit_url.replace("master/MASTG/0x", "master/Document/0x")
+        edit_url = edit_url.replace("master/MASTG/", "master/")
+    elif src_path.startswith("MASWE"):
+        edit_url = f"{edit_url_mastg}{src_path}"
+        edit_url = edit_url.replace("master/MASWE/", "master/weaknesses/")
+    elif src_path.startswith(("contributing", "donate")):
+        edit_url = f"{edit_url_mastg}{src_path}"
+        edit_url = edit_url.replace("master/", "master/docs/")
+    else:
+        edit_url = ""
+    
+    return edit_url
+
+def on_pre_page(page, config, files):
+    try:
+        edit_url_mastg = config["editor_url_mastg"]
+        edit_url_masvs = config["editor_url_masvs"]
+    except KeyError:
+        return page
+    
+    src_path = page.file.src_path
+
+    if src_path.startswith(("MASTG", "MASVS", "MASWE", "contributing", "donate")):
+        edit_url = get_edit_url(src_path, edit_url_mastg, edit_url_masvs)
+        if edit_url.endswith("/index.md"):
+            page.edit_url = ""
+        else:
+            page.edit_url = edit_url
+    else:
+        page.edit_url = ""
+ 
+    return page
\ No newline at end of file
diff --git a/mkdocs.yml b/mkdocs.yml
index 4befe8a69e..6c778c4702 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -2,7 +2,8 @@ site_name: OWASP Mobile Application Security
 repo_url: https://github.com/OWASP/owasp-mastg
 repo_name: OWASP/owasp-mastg
 # use_directory_urls: false # only set for mkdocs build
-edit_uri: "" # disable edit button
+editor_url_mastg: "https://github.com/OWASP/owasp-mastg/edit/master/"
+editor_url_masvs: "https://github.com/OWASP/owasp-masvs/edit/master/"
 
 nav:
   - Home: index.md
@@ -251,6 +252,7 @@ theme:
     - navigation.tracking
     - navigation.indexes
     - content.code.copy
+    - content.action.edit
 
   palette:
     # - primary: #499FFF 
@@ -321,6 +323,7 @@ hooks:
   - docs/hooks/maswe-beta-banner.py
   - docs/hooks/add-tags.py
   - docs/hooks/add-cross-references.py
+  - docs/hooks/edit_button_uri.py
 
 extra:
   generator: false # removed but recreated in copyright above