101 lines
3.3 KiB
Python
101 lines
3.3 KiB
Python
import os
|
|
import smtplib
|
|
import yaml
|
|
import psutil
|
|
import threading
|
|
import time
|
|
from flask import Flask, render_template
|
|
from email.mime.text import MIMEText
|
|
import subprocess
|
|
|
|
app = Flask(__name__)
|
|
|
|
# Load configuration
|
|
try:
|
|
with open("config.yaml") as f:
|
|
config = yaml.safe_load(f)
|
|
print("[DEBUG] Loaded config.yaml successfully.")
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to load config.yaml: {e}")
|
|
exit(1)
|
|
|
|
results = {}
|
|
|
|
def get_folder_size(path):
|
|
print(f"[DEBUG] Calculating size for: {path}")
|
|
try:
|
|
out = subprocess.check_output(["du", "-sb", path], timeout=60)
|
|
size_bytes = int(out.decode().split()[0])
|
|
print(f"[DEBUG] Total size for {path}: {size_bytes / 1e9:.2f} GB")
|
|
return size_bytes
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to get size for {path}: {e}")
|
|
return 0
|
|
|
|
def send_alert(folder, used, limit):
|
|
print(f"[ALERT] Sending email for {folder}: used={used}, limit={limit}")
|
|
try:
|
|
msg = MIMEText(f"Folder {folder} is over limit: {used / 1e9:.2f} GB used of {limit / 1e9:.2f} GB")
|
|
msg['Subject'] = f"[ALERT] Disk usage for {folder}"
|
|
msg['From'] = config['email']['from']
|
|
msg['To'] = config['email']['to']
|
|
|
|
with smtplib.SMTP(config['email']['smtp_server'], config['email']['smtp_port']) as s:
|
|
s.starttls()
|
|
s.login(config['email']['username'], config['email']['password'])
|
|
s.send_message(msg)
|
|
print(f"[INFO] Email alert sent for {folder}")
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to send email alert: {e}")
|
|
|
|
def check_folders():
|
|
print("[INFO] Running folder usage check...")
|
|
new_results = {}
|
|
for folder, limit_str in config['folders'].items():
|
|
if not os.path.exists(folder):
|
|
print(f"[WARN] Folder does not exist: {folder}")
|
|
continue
|
|
try:
|
|
limit = int(limit_str.replace("GB", "").strip()) * 1024 ** 3
|
|
used = get_folder_size(folder)
|
|
exceeded = used > limit
|
|
new_results[folder] = {
|
|
"used_bytes": used,
|
|
"limit_bytes": limit,
|
|
"used_gb": used / (1024 ** 3),
|
|
"limit_gb": limit / (1024 ** 3),
|
|
"exceeded": exceeded
|
|
}
|
|
if exceeded:
|
|
send_alert(folder, used, limit)
|
|
except Exception as e:
|
|
print(f"[ERROR] Failed to check folder {folder}: {e}")
|
|
results.clear()
|
|
results.update(new_results)
|
|
|
|
def schedule_check(interval_sec):
|
|
def loop():
|
|
while True:
|
|
check_folders()
|
|
time.sleep(interval_sec)
|
|
t = threading.Thread(target=loop, daemon=True)
|
|
t.start()
|
|
|
|
@app.route("/")
|
|
def index():
|
|
import datetime
|
|
return render_template("index.html", results=results, now=datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
interval = config.get("schedule", {}).get("seconds", 86400) # default 5 minutes
|
|
print(f"[INFO] Starting scheduled checks every {interval} seconds.")
|
|
schedule_check(interval)
|
|
|
|
if config.get("webui", {}).get("enabled", False):
|
|
port = config["webui"].get("port", 8080)
|
|
print(f"[INFO] Starting Flask Web UI on port {port}")
|
|
app.run(host="0.0.0.0", port=port)
|
|
else:
|
|
print("[INFO] Web UI disabled or missing from config.")
|