Created
January 19, 2025 07:08
-
-
Save madhur/436481bf712866ed57af7007552912f9 to your computer and use it in GitHub Desktop.
Disk monitor script
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/env python3 | |
import os | |
import psutil | |
import json | |
import smtplib | |
from email.message import EmailMessage | |
from datetime import datetime | |
import logging | |
from pathlib import Path | |
import matplotlib.pyplot as plt | |
import matplotlib.dates as mdates | |
from datetime import datetime, timedelta | |
import csv | |
class DiskUsageMonitor: | |
def __init__(self, config_path='config.json'): | |
self.logger = self._setup_logging() | |
self.config = self._load_config(config_path) | |
self.history_file = 'disk_usage_history.csv' | |
self.ensure_history_file() | |
def _setup_logging(self): | |
logging.basicConfig( | |
level=logging.INFO, | |
format='%(asctime)s - %(levelname)s - %(message)s', | |
handlers=[ | |
logging.FileHandler('disk_monitor.log'), | |
logging.StreamHandler() | |
] | |
) | |
return logging.getLogger(__name__) | |
def _load_config(self, config_path): | |
default_config = { | |
'services': { | |
'mysql': { | |
'paths': ['/home/madhur/Downloads/mysql57'], | |
'threshold_gb': 10 | |
}, | |
'kafka': { | |
'paths': ['/home/madhur/kafka'], | |
'threshold_gb': 20 | |
}, | |
'zookeeper': { | |
'paths': ['/home/madhur/zookeeper'], | |
'threshold_gb': 5 | |
}, | |
'prometheus': { | |
'paths': ['/home/madhur/prometheus'], | |
'threshold_gb': 15 | |
} | |
}, | |
'email': { | |
'enabled': True, | |
'smtp_server': 'mail.madhur.co.in', | |
'smtp_port': 587, | |
'sender': 'ahuja.madhur@gmail.com', | |
'recipients': ['ahuja.madhur@gmail.com', 'madhur@mailcow.madhur.co.in'] | |
} | |
} | |
try: | |
with open(config_path, 'r') as f: | |
config = json.load(f) | |
# Merge with default config | |
for service in default_config['services']: | |
if service not in config.get('services', {}): | |
config.setdefault('services', {})[service] = default_config['services'][service] | |
if 'email' not in config: | |
config['email'] = default_config['email'] | |
return config | |
except FileNotFoundError: | |
self.logger.warning(f"Config file {config_path} not found. Using default configuration.") | |
return default_config | |
def get_directory_size(self, path): | |
"""Calculate total size of a directory in bytes.""" | |
try: | |
total_size = 0 | |
for dirpath, _, filenames in os.walk(path): | |
for f in filenames: | |
fp = os.path.join(dirpath, f) | |
if not os.path.islink(fp): # Skip symbolic links | |
total_size += os.path.getsize(fp) | |
return total_size | |
except Exception as e: | |
self.logger.error(f"Error calculating size for {path}: {str(e)}") | |
return 0 | |
def ensure_history_file(self): | |
"""Create history CSV file if it doesn't exist.""" | |
if not os.path.exists(self.history_file): | |
with open(self.history_file, 'w', newline='') as f: | |
writer = csv.writer(f) | |
headers = ['timestamp'] + list(self.config['services'].keys()) | |
writer.writerow(headers) | |
def save_to_history(self, usage_data): | |
"""Save current usage data to history CSV.""" | |
timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S') | |
row = [timestamp] + [usage_data[service]['size'] for service in self.config['services'].keys()] | |
with open(self.history_file, 'a', newline='') as f: | |
writer = csv.writer(f) | |
writer.writerow(row) | |
def generate_usage_graph(self): | |
"""Generate a graph showing disk usage trends.""" | |
data = [] | |
with open(self.history_file, 'r') as f: | |
reader = csv.reader(f) | |
headers = next(reader) | |
for row in reader: | |
data.append(row) | |
if not data: | |
return None | |
# Convert data for plotting | |
timestamps = [datetime.strptime(row[0], '%Y-%m-%d %H:%M:%S') for row in data] | |
services = headers[1:] | |
# Create the plot | |
plt.figure(figsize=(12, 6)) | |
for i, service in enumerate(services): | |
usage_values = [float(row[i+1]) for row in data] | |
plt.plot(timestamps, usage_values, marker='o', label=service) | |
plt.title('Disk Usage Trends') | |
plt.xlabel('Time') | |
plt.ylabel('Usage (GB)') | |
plt.grid(True) | |
plt.legend() | |
# Format x-axis | |
plt.gcf().autofmt_xdate() | |
plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d %H:%M')) | |
# Save the plot | |
plot_path = 'disk_usage_trend.png' | |
plt.savefig(plot_path) | |
plt.close() | |
return plot_path | |
def check_disk_usage(self): | |
"""Check disk usage for all configured services.""" | |
alerts = [] | |
usage_data = {} | |
for service, config in self.config['services'].items(): | |
total_size_bytes = 0 | |
for path in config['paths']: | |
if os.path.exists(path): | |
size = self.get_directory_size(path) | |
total_size_bytes += size | |
else: | |
self.logger.warning(f"Path {path} for service {service} does not exist") | |
total_size_gb = total_size_bytes / (1024**3) # Convert to GB | |
threshold_gb = config['threshold_gb'] | |
usage_data[service] = { | |
'size': total_size_gb, | |
'threshold': threshold_gb | |
} | |
if total_size_gb > threshold_gb: | |
message = (f"Alert: {service} is using {total_size_gb:.2f}GB, " | |
f"exceeding threshold of {threshold_gb}GB") | |
alerts.append(message) | |
self.logger.warning(message) | |
else: | |
self.logger.info(f"{service} usage: {total_size_gb:.2f}GB (Threshold: {threshold_gb}GB)") | |
# Save current usage data to history | |
self.save_to_history(usage_data) | |
return alerts, usage_data | |
def send_email_alert(self, alerts, usage_data): | |
"""Send email alerts if configured.""" | |
if not self.config['email']['enabled'] or not alerts: | |
return | |
email_config = self.config['email'] | |
# Generate the plot | |
plot_path = self.generate_usage_graph() | |
# Create email content | |
content = [] | |
content.append("Disk Usage Alerts:\n") | |
content.extend(alerts) | |
content.append("\n\nCurrent Usage Summary:") | |
for service, data in usage_data.items(): | |
content.append(f"\n{service}:") | |
content.append(f" Current Usage: {data['size']:.2f}GB") | |
content.append(f" Threshold: {data['threshold']}GB") | |
msg = EmailMessage() | |
msg.set_content('\n'.join(content)) | |
# Add the plot as attachment if it was generated | |
if plot_path and os.path.exists(plot_path): | |
with open(plot_path, 'rb') as f: | |
image_data = f.read() | |
msg.add_attachment(image_data, maintype='image', subtype='png', filename='disk_usage_trend.png') | |
msg['Subject'] = 'Disk Usage Alert' | |
msg['From'] = email_config['sender'] | |
msg['To'] = ', '.join(email_config['recipients']) | |
try: | |
with smtplib.SMTP(email_config['smtp_server'], email_config['smtp_port']) as server: | |
server.starttls() | |
server.send_message(msg) | |
self.logger.info("Email alert sent successfully") | |
except Exception as e: | |
self.logger.error(f"Failed to send email alert: {str(e)}") | |
def main(): | |
monitor = DiskUsageMonitor() | |
alerts, usage_data = monitor.check_disk_usage() | |
if alerts: | |
monitor.send_email_alert(alerts, usage_data) | |
# Generate and save the plot even if there are no alerts | |
monitor.generate_usage_graph() | |
if __name__ == "__main__": | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment