307 lines
11 KiB
Python
307 lines
11 KiB
Python
|
|
###############################################################
|
|
# SSDObject.py
|
|
# ssd object class function created for ssd monitor
|
|
# functions for handling database and shell commands
|
|
# class definitions and other class helper functions
|
|
###############################################################
|
|
|
|
import sqlite3
|
|
import json, redis
|
|
import subprocess
|
|
|
|
db_path = '/opt/ssd_health/drive_records.db'
|
|
debug_output = False
|
|
show_records = False
|
|
suppress_errors = False
|
|
push_redis = False
|
|
|
|
###############################################################
|
|
# Class Definition
|
|
#
|
|
# Most of the heavy lifting is done on object instantiation
|
|
# Object is instantiated with just the dev_id and flavor
|
|
# Drive data is collected on instantition and stored to database
|
|
###############################################################
|
|
|
|
class SSDObject:
|
|
def __str__(self) -> str:
|
|
return f"""Drive at /dev/{self.dev_id} is a '{self.model}' {self.capacity} {self.flavor}, SN: {self.serial}"""
|
|
|
|
def __init__(self, dev_id: str):
|
|
self.dev_id = dev_id
|
|
if not check_serial_attached(self.dev_id):
|
|
raise TypeError(f"No device at /dev/{self.dev_id}")
|
|
self._smart_data = return_smartctl(dev_id)
|
|
self.flavor = self._get_flavor()
|
|
if self.flavor == "HDD":
|
|
raise TypeError("Unable to instantiate HDD")
|
|
if self.flavor == "Error":
|
|
raise TypeError("Unable to instantiate storage device")
|
|
self.serial = self._smart_data['serial_number']
|
|
self.model = self._smart_data['model_name']
|
|
self.capacity_bytes = self._smart_data['user_capacity']['bytes']
|
|
self.smart_status = self._smart_data['smart_status']['passed']
|
|
self.capacity = self._get_human_capacity()
|
|
self.sector_size = return_sector_size(self.dev_id)
|
|
self.gb_written = self._get_gbw()
|
|
self._update_db()
|
|
|
|
def _get_flavor(self) -> str:
|
|
if "rotation_rate" in json.dumps(self._smart_data):
|
|
if int(self._smart_data['rotation_rate'] == 0):
|
|
if not suppress_errors:
|
|
print(f"Warning - /dev/{self.dev_id} is a weird SSD with a rotation rate of 0")
|
|
return "SSD"
|
|
else:
|
|
return "HDD"
|
|
elif "NVMe" in json.dumps(self._smart_data):
|
|
return "NVMe"
|
|
elif "Solid State" in json.dumps(self._smart_data):
|
|
return "SSD"
|
|
elif "Unknown USB bridge" in json.dumps(self._smart_data):
|
|
return "Error"
|
|
|
|
def _get_human_capacity(self) -> str:
|
|
size = self.capacity_bytes
|
|
factor = 1024
|
|
units = [
|
|
(factor ** 4, "TiB"),
|
|
(factor ** 3, "GiB"),
|
|
(factor ** 2, "MiB"),
|
|
(factor ** 1, "KiB"),
|
|
]
|
|
for thresh, suffix in units:
|
|
if size > thresh:
|
|
value = size / thresh
|
|
return f"{value:.{0}f} {suffix}"
|
|
|
|
def _get_gbw(self) -> str:
|
|
result = ''
|
|
gib_factor = 2 ** 30
|
|
if self.flavor == "SSD":
|
|
data_units_written = return_ls_written(self._smart_data)
|
|
result = round(data_units_written * self.sector_size / gib_factor, 2)
|
|
elif self.flavor == "NVMe":
|
|
data_units_written = float(self._smart_data['nvme_smart_health_information_log']['data_units_written'])
|
|
result = round(data_units_written * self.sector_size / gib_factor, 2)
|
|
return result
|
|
|
|
def _update_db(self):
|
|
if push_redis:
|
|
update_disk_redis()
|
|
if self.exists():
|
|
drive_query = f"""
|
|
UPDATE drive_records SET gb_written = '{self.gb_written}', smart = '{self.smart_status}' WHERE serial = '{self.serial}';
|
|
"""
|
|
else:
|
|
drive_query = f"""
|
|
INSERT INTO drive_records (serial, model, flavor, capacity, gb_written, smart)
|
|
VALUES ('{self.serial}', '{self.model}', '{self.flavor}', '{self.capacity}', '{self.gb_written}', '{self.smart_status}');
|
|
|
|
"""
|
|
query_db(drive_query)
|
|
|
|
def exists(self) -> bool:
|
|
return check_serial_exists(self.serial)
|
|
|
|
def attached(self) -> bool:
|
|
return(check_serial_attached(self.serial))
|
|
|
|
|
|
########################################
|
|
# Other Helper Functions
|
|
########################################
|
|
|
|
# subroutine to run a command, return stdout as array unless zero_only then return [0]
|
|
def run_command(cmd, zero_only=False):
|
|
# Run the command and capture the output
|
|
result = subprocess.run(cmd, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
# Decode the byte output to a string
|
|
output = result.stdout.decode('utf-8')
|
|
# Split the output into lines and store it in an array
|
|
output_lines = [line for line in output.split('\n') if line]
|
|
# Return result
|
|
try:
|
|
return output_lines[0] if zero_only else output_lines
|
|
except:
|
|
return output_lines
|
|
|
|
def return_smartctl(drive_id):
|
|
smartctl_string = f"/usr/sbin/smartctl --json -x /dev/{drive_id} || true"
|
|
smartctl_result = subprocess.run(smartctl_string, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
smartctl_data = json.loads(smartctl_result.stdout.decode("utf-8"))
|
|
return smartctl_data
|
|
|
|
def return_sector_size(drive_id):
|
|
sector_size_command = f"fdisk -l /dev/{drive_id} | grep 'Sector size' | awk '{{print $4}}'"
|
|
sector_size_result = subprocess.run(sector_size_command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
return int(sector_size_result.stdout.decode("utf-8"))
|
|
|
|
def return_ls_written(data):
|
|
pages = data.get("ata_device_statistics", {}).get("pages", [])
|
|
for page in pages:
|
|
for entry in page.get("table", []):
|
|
if entry.get("name") == "Logical Sectors Written":
|
|
return entry.get("value")
|
|
|
|
# Function to return all drive records in database
|
|
def get_all_drive_records():
|
|
get_all_drives = "SELECT * FROM drive_records"
|
|
rows = query_db(get_all_drives)
|
|
drives = []
|
|
for row in rows:
|
|
drive = {
|
|
'id': row[0],
|
|
'serial': row[1],
|
|
'model': row[2],
|
|
'flavor': row[3],
|
|
'capacity': row[4],
|
|
'gb_written': round(float(row[5]),2),
|
|
'smart': row[6]
|
|
}
|
|
drives.append(drive)
|
|
return json.dumps(drives)
|
|
|
|
|
|
# return attached disks
|
|
def list_disk_and_serial():
|
|
# Init blank devices array
|
|
devices = []
|
|
# get the devices
|
|
cmd = "lsblk -o NAME,SERIAL,SIZE,TYPE | grep sd | grep disk | awk '{print $1 \",\" $2. \",\" $3}'"
|
|
# try to run the command, should not fail
|
|
try:
|
|
while 'disk' in run_command(cmd, zero_only=False):
|
|
time.sleep(0.5)
|
|
devices = run_command(cmd, zero_only=False)
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"An error occurred: {e.stderr.decode('utf-8')}")
|
|
drives = []
|
|
for device in devices:
|
|
if debug_output:
|
|
print(device)
|
|
drive = {
|
|
"dev_id": device.split(',')[0],
|
|
"serial": device.split(',')[1],
|
|
"capacity": device.split(',')[2]
|
|
}
|
|
drives.append(drive)
|
|
# return the devices as an array
|
|
if debug_output:
|
|
print(drives)
|
|
return drives
|
|
|
|
# Function to check if a serial number exists in the database
|
|
def check_serial_exists(serial):
|
|
serial_check = f"SELECT * FROM drive_records WHERE serial='{serial}'"
|
|
if debug_output:
|
|
print(serial_check)
|
|
return bool(query_db(serial_check))
|
|
|
|
def check_serial_attached(serial):
|
|
serial_check = f"lsblk -o NAME,SERIAL,SIZE,TYPE | grep {serial} || true"
|
|
if run_command(serial_check, zero_only=False):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
# Function to run SQL Query
|
|
def query_db(sql_query):
|
|
try:
|
|
with sqlite3.connect(db_path) as conn:
|
|
cursor = conn.cursor()
|
|
if debug_output:
|
|
print("Executing SQL query:", sql_query)
|
|
cursor.execute(sql_query)
|
|
rows = cursor.fetchall()
|
|
if debug_output:
|
|
print("Query Result:", rows)
|
|
conn.commit()
|
|
conn.close()
|
|
return rows
|
|
except sqlite3.Error as e:
|
|
if not suppress_errors:
|
|
print("An error occurred:", e)
|
|
return []
|
|
|
|
def init_db():
|
|
print("Checking Database...")
|
|
db_check = "SELECT name FROM sqlite_master WHERE type='table' AND name='drive_records';"
|
|
create_table_command = """
|
|
CREATE TABLE drive_records (
|
|
id INTEGER PRIMARY KEY,
|
|
serial TEXT NOT NULL,
|
|
model TEXT NOT NULL,
|
|
flavor TEXT NOT NULL,
|
|
capacity TEXT NOT NULL,
|
|
gb_written TEXT NOT NULL,
|
|
smart TEXT NOT NULL
|
|
);
|
|
"""
|
|
|
|
# check for drives
|
|
try:
|
|
result = bool(query_db(db_check))
|
|
if debug_output:
|
|
print(f"Database exists: {result}")
|
|
# Check if any tables were found
|
|
if result:
|
|
print("drive_records exists, skipping db init")
|
|
if debug_output or show_records:
|
|
all_drives = json.loads(get_all_drive_records())
|
|
print("--- Drive Records ---")
|
|
for drive in all_drives:
|
|
print(f"{drive['model']} - SN: {drive['serial']}")
|
|
print("--- End Records ---")
|
|
print()
|
|
else:
|
|
print("drive_records does not exist, creating")
|
|
try:
|
|
result_init = query_db(create_table_command)
|
|
if debug_output:
|
|
print(result_init)
|
|
print("Database created - 201")
|
|
except sqlite3.Error as e:
|
|
if not suppress_errors:
|
|
print(f"error during table initialization: {e}")
|
|
return jsonify({'error during table initialization - 401': e}), 401
|
|
|
|
except sqlite3.Error as e:
|
|
if not suppress_errors:
|
|
print(f"error during table check: {e}")
|
|
return jsonify({'error during table check - 400': e}), 400
|
|
|
|
|
|
####################################################
|
|
### Redis Functions
|
|
####################################################
|
|
|
|
r = redis.Redis(host='172.17.0.1', port=6379)
|
|
|
|
def update_disk_redis():
|
|
active = list_disk_and_serial()
|
|
all_rec = json.loads(get_all_drive_records())
|
|
enriched = merge_active_with_details(active, all_rec)
|
|
r.publish('attached_disks', json.dumps(enriched))
|
|
if debug_output:
|
|
print("=== Active drives sent to Redis ===")
|
|
print(json.dumps(enriched, indent=2))
|
|
|
|
def merge_active_with_details(active, all_records):
|
|
# Build a quick lookup dictionary keyed by serial
|
|
record_by_serial = {rec['serial']: rec for rec in all_records}
|
|
# Add the extra fields to each active drive
|
|
for drive in active:
|
|
rec = record_by_serial.get(drive['serial'])
|
|
if rec:
|
|
extra = {k: v for k, v in rec.items() if k not in ('id', 'serial')}
|
|
drive.update(extra)
|
|
return active
|
|
|
|
########################################
|
|
# Run init_db when Class file is imported
|
|
########################################
|
|
|
|
init_db()
|