Add disk reporting with a feature flag

README.md

@@ -58,6 +58,7 @@ pyninja start
 - **NINJA_PORT** - Port number for the API server.
 - **REMOTE_EXECUTION** - Boolean flag to enable remote execution.
 - **API_SECRET** - Secret access key for running commands on server remotely.
+- **DISK_REPORT** - Boolean flag to enable disk report feature using [PyUdisk].
 - **MONITOR_USERNAME** - Username to authenticate the monitoring page.
 - **MONITOR_PASSWORD** - Password to authenticate the monitoring page.
 - **MONITOR_SESSION** - Session timeout for the monitoring page.
@@ -153,3 +154,4 @@ Licensed under the [MIT License][license]
 [license]: https://github.com/thevickypedia/PyNinja/blob/master/LICENSE
 [runbook]: https://thevickypedia.github.io/PyNinja/
 [samples]: https://github.com/thevickypedia/PyNinja/tree/main/samples
+[PyUdisk]: https://github.com/thevickypedia/PyUdisk

docs/genindex.html

@@ -112,11 +112,13 @@ <h2 id="A">A</h2>
 </li>
       <li><a href="index.html#pyninja.modules.models.EnvConfig.api_secret">api_secret (pyninja.modules.models.EnvConfig attribute)</a>
 </li>
-  </ul></td>
-  <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="index.html#pyninja.modules.models.EnvConfig.apikey">apikey (pyninja.modules.models.EnvConfig attribute)</a>
 </li>
+  </ul></td>
+  <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="index.html#pyninja.modules.exceptions.APIResponse">APIResponse</a>
+</li>
+      <li><a href="index.html#pyninja.modules.models.EnvConfig.assert_disk_report">assert_disk_report() (pyninja.modules.models.EnvConfig class method)</a>
 </li>
       <li><a href="index.html#pyninja.modules.models.Session.auth_counter">auth_counter (pyninja.modules.models.Session attribute)</a>
 </li>
@@ -178,15 +180,19 @@ <h2 id="D">D</h2>
 </li>
       <li><a href="index.html#pyninja.modules.payloads.ListFiles.deep_scan">deep_scan (pyninja.modules.payloads.ListFiles attribute)</a>
 </li>
-  </ul></td>
-  <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="index.html#pyninja.features.operations.default">default() (in module pyninja.features.operations)</a>
 </li>
+  </ul></td>
+  <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="index.html#pyninja.modules.models.ServiceStatus.description">description (pyninja.modules.models.ServiceStatus attribute)</a>
 </li>
       <li><a href="index.html#pyninja.modules.payloads.ListFiles.directory">directory (pyninja.modules.payloads.ListFiles attribute)</a>
 </li>
       <li><a href="index.html#pyninja.modules.models.EnvConfig.disk_lib">disk_lib (pyninja.modules.models.EnvConfig attribute)</a>
+</li>
+      <li><a href="index.html#pyninja.modules.models.EnvConfig.disk_report">disk_report (pyninja.modules.models.EnvConfig attribute)</a>
+</li>
+      <li><a href="index.html#pyninja.monitor.resources.disk_utils_metrics">disk_utils_metrics() (in module pyninja.monitor.resources)</a>
 </li>
       <li><a href="index.html#pyninja.modules.models.DiskLib">DiskLib (class in pyninja.modules.models)</a>
 </li>

docs/index.html

@@ -1722,6 +1722,11 @@ <h2>Models<a class="headerlink" href="#models" title="Permalink to this heading"
 <span class="sig-name descname"><span class="pre">monitor_session</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">int</span></em><a class="headerlink" href="#pyninja.modules.models.EnvConfig.monitor_session" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
+<dl class="py attribute">
+<dt class="sig sig-object py" id="pyninja.modules.models.EnvConfig.disk_report">
+<span class="sig-name descname"><span class="pre">disk_report</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">bool</span></em><a class="headerlink" href="#pyninja.modules.models.EnvConfig.disk_report" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
 <dl class="py attribute">
 <dt class="sig sig-object py" id="pyninja.modules.models.EnvConfig.max_connections">
 <span class="sig-name descname"><span class="pre">max_connections</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">int</span></em><a class="headerlink" href="#pyninja.modules.models.EnvConfig.max_connections" title="Permalink to this definition">¶</a></dt>
@@ -1777,6 +1782,17 @@ <h2>Models<a class="headerlink" href="#models" title="Permalink to this heading"
 <span class="sig-name descname"><span class="pre">log_config</span></span><em class="property"><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="pre">Optional</span><span class="p"><span class="pre">[</span></span><span class="pre">Union</span><span class="p"><span class="pre">[</span></span><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Any</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">Path</span><span class="p"><span class="pre">]</span></span><span class="p"><span class="pre">]</span></span></em><a class="headerlink" href="#pyninja.modules.models.EnvConfig.log_config" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
+<dl class="py method">
+<dt class="sig sig-object py" id="pyninja.modules.models.EnvConfig.assert_disk_report">
+<em class="property"><span class="pre">classmethod</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">assert_disk_report</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">value</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">bool</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">bool</span></span></span><a class="headerlink" href="#pyninja.modules.models.EnvConfig.assert_disk_report" title="Permalink to this definition">¶</a></dt>
+<dd><p>Ensure disk_report is enabled only for Linux machines.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters<span class="colon">:</span></dt>
+<dd class="field-odd"><p><strong>value</strong> – Takes the user input as an argument.</p>
+</dd>
+</dl>
+</dd></dl>
+
 <dl class="py method">
 <dt class="sig sig-object py" id="pyninja.modules.models.EnvConfig.parse_api_secret">
 <em class="property"><span class="pre">classmethod</span><span class="w"> </span></em><span class="sig-name descname"><span class="pre">parse_api_secret</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">value</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">str</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">str</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">None</span></span></span><a class="headerlink" href="#pyninja.modules.models.EnvConfig.parse_api_secret" title="Permalink to this definition">¶</a></dt>
@@ -2367,6 +2383,12 @@ <h1>PyNinja - Monitor<a class="headerlink" href="#pyninja-monitor" title="Permal
 </dl>
 </dd></dl>
 
+<dl class="py function">
+<dt class="sig sig-object py" id="pyninja.monitor.resources.disk_utils_metrics">
+<em class="property"><span class="k"><span class="pre">async</span></span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">pyninja.monitor.resources.</span></span><span class="sig-name descname"><span class="pre">disk_utils_metrics</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">dict</span><span class="p"><span class="pre">]</span></span></span></span><a class="headerlink" href="#pyninja.monitor.resources.disk_utils_metrics" title="Permalink to this definition">¶</a></dt>
+<dd><p>Placeholder for future metrics.</p>
+</dd></dl>
+
 <dl class="py function">
 <dt class="sig sig-object py" id="pyninja.monitor.resources.system_resources">
 <em class="property"><span class="k"><span class="pre">async</span></span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">pyninja.monitor.resources.</span></span><span class="sig-name descname"><span class="pre">system_resources</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">dict</span><span class="p"><span class="pre">]</span></span></span></span><a class="headerlink" href="#pyninja.monitor.resources.system_resources" title="Permalink to this definition">¶</a></dt>

pyninja/modules/models.py

@@ -190,6 +190,8 @@ class EnvConfig(BaseSettings):
     monitor_username: str | None = None
     monitor_password: str | None = None
     monitor_session: PositiveInt = 3_600
+    # todo: disk_lib env var will contradict with PyUdisk
+    disk_report: bool = False
     max_connections: PositiveInt = 3
     no_auth: bool = False
     processes: List[str] = []
@@ -202,6 +204,28 @@ class EnvConfig(BaseSettings):
     rate_limit: RateLimit | List[RateLimit] = []
     log_config: Dict[str, Any] | FilePath | None = None
 
+    # noinspection PyMethodParameters
+    @field_validator("disk_report", mode="after")
+    def assert_disk_report(cls, value: bool) -> bool:
+        """Ensure disk_report is enabled only for Linux machines.
+
+        Args:
+            value: Takes the user input as an argument.
+        """
+        if value:
+            assert OPERATING_SYSTEM == "linux", ValueError(
+                "disk_report feature can be enabled only on Linux machines!"
+            )
+            try:
+                from pyudisk.config import EnvConfig as PyUdiskConfig
+            except (ImportError, ModuleNotFoundError):
+                raise ValueError(
+                    "PyUdisk has not been installed. Use pip install 'PyNinja[extra]' to use disk reporting feature."
+                )
+            disk_lib = PyUdiskConfig().disk_lib
+            assert os.path.exists(disk_lib), ValueError(f"{disk_lib!r} path does not exist!")
+        return value
+
     # noinspection PyMethodParameters
     @field_validator("api_secret", mode="after")
     def parse_api_secret(cls, value: str | None) -> str | None:

pyninja/monitor/drive.py

@@ -1,6 +1,7 @@
 import logging
 
 from fastapi.responses import HTMLResponse
+from pyninja import monitor
 
 LOGGER = logging.getLogger("uvicorn.default")
 
@@ -15,3 +16,12 @@ async def report() -> HTMLResponse:
     import pyudisk
 
     return HTMLResponse(content=pyudisk.generate_report(raw=True))
+
+
+async def invalid(ctx: str, status_code: int = 500):
+    return await monitor.config.clear_session(
+        response=HTMLResponse(
+            content=ctx, status_code=status_code
+        )
+    )
+

pyninja/monitor/resources.py

@@ -182,13 +182,27 @@ async def get_system_metrics() -> Dict[str, dict]:
     )
 
 
+async def disk_utils_metrics() -> Dict[str, dict]:
+    """Placeholder for future metrics."""
+    from pyudisk.main import EnvConfig, smart_metrics
+
+    for disk in smart_metrics(EnvConfig()):
+        print(disk.Info.Model)
+        print(disk.Partition.IdLabel)
+        print(disk.Partition.MountPoints)
+        print(disk.Usage)
+        print(disk.Attributes.SmartTemperature)
+        print(disk.Attributes.SmartUpdated)
+
+
 async def system_resources() -> Dict[str, dict]:
     """Gather system resources including Docker stats asynchronously.
 
     Returns:
         Dict[str, dict]:
         Returns a nested dictionary.
     """
+    # await disk_utils_metrics()
     system_metrics_task = asyncio.create_task(get_system_metrics())
     docker_stats_task = asyncio.create_task(get_docker_stats())
     service_stats_task = asyncio.create_task(

pyninja/monitor/routes.py

@@ -132,17 +132,16 @@ async def monitor_endpoint(
                 name="main.html", context=ctx
             )
         elif render == "drive":
-            LOGGER.info("Rendering disk report!")
-            try:
-                response = await monitor.drive.report()
-                # If drive option is chosen during login page, the cookie is deleted once logged in!
-            except Exception as error:
-                LOGGER.error(error)
-                response = await monitor.config.clear_session(
-                    response=HTMLResponse(
-                        content="Failed to generate disk report", status_code=500
-                    )
-                )
+            if models.env.disk_report:
+                LOGGER.info("Rendering disk report!")
+                try:
+                    response = await monitor.drive.report()
+                    # If drive option is chosen during login page, the cookie is deleted once logged in!
+                except Exception as error:
+                    LOGGER.error(error)
+                    response = await monitor.drive.invalid("Failed to generate disk report")
+            else:
+                response = await monitor.drive.invalid("Disk reporting feature is not enabled in the server!")
             response.delete_cookie(key="render")
             return response
         # todo: Implement a better way to handle this
@@ -154,8 +153,7 @@ async def monitor_endpoint(
             "request": request,
             "signin": "/login",
             "version": f"v{version.__version__}",
-            # todo: Add an env var to enable reporting manually
-            "linux": models.OPERATING_SYSTEM == "linux",
+            "disk_report": models.env.disk_report,
         },
     )
 

pyninja/monitor/templates/index.html

@@ -150,7 +150,7 @@ <h2 style="margin-top:5%">This page requires JavaScript
                 <i class="fa-regular fa-eye" id="eye"></i>
             </div>
             <button type="submit" onclick="submitToAPI(event, 'monitor')">Sign In</button>
-            {% if linux %}
+            {% if disk_report %}
                 <br>
                 <button type="submit" onclick="submitToAPI(event, 'drive')">Disk Report</button>
             {% endif %}