From f2aabf6a1da5ce1387747f9c0e54bcda4e0fb0e9 Mon Sep 17 00:00:00 2001
From: BennyThink <benny.think@gmail.com>
Date: Fri, 20 Aug 2021 23:48:56 +0800
Subject: [PATCH] add vnstat

---
 Dockerfile | 10 ++++++----
 limit.py   | 14 +++++++++++++-
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index e895c6e..4828343 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -6,12 +6,14 @@ RUN pip3 install --user -r /tmp/requirements.txt && rm /tmp/requirements.txt
 
 
 FROM python:alpine
-RUN apk update && apk add  --no-cache ffmpeg
+WORKDIR /ytdlbot
+ENV TZ=Asia/Shanghai
+
+RUN apk update && apk add  --no-cache ffmpeg vnstat
 COPY --from=builder /root/.local /usr/local
 COPY --from=builder /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/
 COPY --from=builder /usr/share/zoneinfo /usr/share/zoneinfo
 COPY . /ytdlbot
+RUN echo "/usr/sbin/vnstatd -d;/usr/local/bin/python ytdl.py"> /ytdlbot/start.sh
 
-WORKDIR /ytdlbot
-ENV TZ=Asia/Shanghai
-CMD ["python", "ytdl.py"]
\ No newline at end of file
+CMD ["sh", "start.sh"]
\ No newline at end of file
diff --git a/limit.py b/limit.py
index edd3290..1d35f44 100644
--- a/limit.py
+++ b/limit.py
@@ -13,6 +13,7 @@ import math
 import os
 import re
 import sqlite3
+import subprocess
 import tempfile
 import time
 from io import BytesIO
@@ -39,6 +40,7 @@ class Redis:
         quota_banner = "=" * 20 + "Quota" + "=" * 20
         metrics_banner = "=" * 20 + "Metrics" + "=" * 20
         usage_banner = "=" * 20 + "Usage" + "=" * 20
+        vnstat_banner = "=" * 20 + "vnstat" + "=" * 20
         self.final_text = f"""
 {db_banner}
 %s
@@ -53,6 +55,10 @@ class Redis:
 
 
 {usage_banner}
+%s
+
+
+{vnstat_banner}
 %s
         """
 
@@ -108,7 +114,13 @@ class Redis:
         fd.sort(key=lambda x: int(x[1]))
         quota_text = self.generate_table(["UserID", "bytes", "human readable", "refresh time"], fd)
 
-        return self.final_text % (db_text, quota_text, metrics_text, usage_text)
+        # vnstat
+        if os.uname().sysname == "Darwin":
+            cmd = "/usr/local/bin/vnstat -i en0".split()
+        else:
+            cmd = "/usr/bin/vnstat -i eth0".split()
+        vnstat_text = subprocess.check_output(cmd).decode('u8')
+        return self.final_text % (db_text, quota_text, metrics_text, usage_text, vnstat_text)
 
     def reset_today(self):
         pairs = self.r.hgetall("metrics")