From 9a7930bc2579a6c6a5808d545e4a77a65a25a4a5 Mon Sep 17 00:00:00 2001
From: dettlaff <dettlaff@riseup.net>
Date: Thu, 25 Jul 2024 20:01:48 +0400
Subject: [PATCH] feat: add net usage

---
 selfprivacy_api/utils/prometheus.py | 40 +++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/selfprivacy_api/utils/prometheus.py b/selfprivacy_api/utils/prometheus.py
index 60f12aa..a8b4daa 100644
--- a/selfprivacy_api/utils/prometheus.py
+++ b/selfprivacy_api/utils/prometheus.py
@@ -146,3 +146,43 @@ class PrometheusQueries:
             end_timestamp,
             step,
         )
+
+    @staticmethod
+    def network_usage(
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        step: int = 60,  # seconds
+    ) -> PrometheusQueryResult:
+        """
+        Get network usage information for both download and upload.
+
+        Args:
+            start (datetime, optional): The start time.
+                Defaults to 20 minutes ago if not provided.
+            end (datetime, optional): The end time.
+                Defaults to current time if not provided.
+            step (int): Interval in seconds for querying network data.
+        """
+
+        if start is None:
+            start = datetime.now() - timedelta(minutes=20)
+
+        if end is None:
+            end = datetime.now()
+
+        start_timestamp = int(start.timestamp())
+        end_timestamp = int(end.timestamp())
+
+        query = """
+        (
+            sum(rate(node_network_receive_bytes_total{device!="lo"}[5m])) as download,
+            sum(rate(node_network_transmit_bytes_total{device!="lo"}[5m])) as upload
+        )
+        """
+
+        return PrometheusQueries._send_query(
+            query,
+            start_timestamp,
+            end_timestamp,
+            step,
+        )