Skip to content

Instantly share code, notes, and snippets.

@madhur
Created January 19, 2025 07:08
Show Gist options
  • Save madhur/436481bf712866ed57af7007552912f9 to your computer and use it in GitHub Desktop.
Save madhur/436481bf712866ed57af7007552912f9 to your computer and use it in GitHub Desktop.
Disk monitor script
#!/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