reHDDPrinter/layouter.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Author:                     Hendrik Schutter, localhorst@mosad.xyz
Date of creation:           2022/11/16
Date of last modification:  2025/06/08
"""

import re
import dataclasses
import glob
import datetime
import json
import logging
import qrcode
from PIL import Image, ImageFont, ImageDraw

# Constants
FONT_PATH = "/usr/share/fonts"
DEFAULT_FONT_REGULAR = "DejaVuSans.ttf"
DEFAULT_FONT_BOLD = "DejaVuSans-Bold.ttf"
OUTPUT_WIDTH = 696  # px
OUTPUT_HEIGHT = 190  # px
TEXT_X_OFFSET = 190  # px
QR_CODE_SIZE = 179  # px
# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)


@dataclasses.dataclass
class DriveData:
    drive_index: int
    drive_state: str
    drive_connection_type: str
    modelfamily: str
    modelname: str
    capacity: int
    serialnumber: str
    power_on_hours: int
    power_cycle: int
    smart_error_count: int
    shred_timestamp: int
    shred_duration: int


@dataclasses.dataclass
class DriveDataJson:
    state: str
    contype: str
    fam: str
    name: str
    cap: int
    sn: str
    poh: int
    pc: int
    err: int
    time: int
    dur: int


@dataclasses.dataclass
class DriveDataPrintable:
    connectiontype: str
    modelfamily: str
    modelname: str
    capacity: str
    serialnumber: str
    power_on_hours: str
    power_cycle: str
    smart_error_count: str
    shred_timestamp: str
    shred_duration: str


@dataclasses.dataclass
class ReHddInfo:
    link: str
    version: str


@dataclasses.dataclass
class QrDataJson:
    drive: DriveDataJson
    rehdd: ReHddInfo


def find_font_path(font_name):
    """Finds the full path of the specified font."""
    try:
        files = glob.glob(f"{FONT_PATH}/**/{font_name}", recursive=True)
        return files[0] if files else None
    except Exception as e:
        logging.error(f"Error locating font {font_name}: {e}")
        return None


def human_readable_capacity(size, decimal_places=0, base=1000):
    """Converts bytes to a human-readable string."""
    units = (
        ["B", "KB", "MB", "GB", "TB", "PB"]
        if base == 1000
        else ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]
    )
    for unit in units:
        if size < base or unit == units[-1]:
            return f"{size:.{decimal_places}f} {unit}"
        size /= base


def human_readable_power_on_hours(hours):
    """Converts hours to human-readable string in hours, days, and years."""
    return (
        str(hours)
        + "h or "
        + str(int(hours / 24))
        + "d or "
        + str("{:.2f}".format(float(hours / 24 / 365)))
        + "y"
    )


def cut_string(max_length, data, direction="end"):
    """Trims a string to the maximum length, adding ellipses if necessary."""
    if len(data) > max_length:
        return (
            f"{data[:max_length-4]} ..."
            if direction == "end"
            else f"... {data[-(max_length-4):]}"
        )
    return data


def format_to_printable(drive):
    return DriveDataPrintable(
        drive.drive_connection_type,
        cut_string(20, re.sub(r"[^a-zA-Z0-9. ]", "", drive.modelfamily), "end"),
        cut_string(20, re.sub(r"[^a-zA-Z0-9. ]", "", drive.modelname), "end"),
        cut_string(20, human_readable_capacity(drive.capacity), "end"),
        cut_string(20, re.sub(r"[^a-zA-Z0-9.-_]", "", drive.serialnumber), "start"),
        cut_string(30, human_readable_power_on_hours(drive.power_on_hours), "end"),
        cut_string(10, str(drive.power_cycle), "end"),
        cut_string(10, str(drive.smart_error_count), "end"),
        cut_string(
            30,
            datetime.datetime.fromtimestamp(
                drive.shred_timestamp, datetime.UTC
            ).strftime("%Y-%m-%d %H:%M:%S"),
            "end",
        ),
        cut_string(30, str(datetime.timedelta(seconds=drive.shred_duration)), "end"),
    )

def draw_text(drawable, printable_data, font_regular, font_bold, font_bold_bigger):
    """Draws formatted text with Cycles and Errors on one row."""
    y_start = 4
    line_height = 26
    label_x = TEXT_X_OFFSET
    value_offset = 115
    right_field_spacing = 200  # Horizontal gap between Cycles and Errors
    x_capacity = 520
    y_capacity = 142
    x_connection_type = 600
    y_connection_type = y_start
    y_spacing_connection_type = 25
    

    # Serial Number
    drawable.text((label_x, y_start), "Serial:", fill=0, font=font_bold)
    drawable.text(
        (label_x + value_offset, y_start),
        printable_data.serialnumber,
        fill=0,
        font=font_bold,
    )

    y1 = y_start + line_height
    y2 = y1 + line_height
    y3 = y2 + line_height
    y4 = y3 + line_height
    y5 = y4 + line_height
    y6 = y5 + line_height

    # Left-Aligned Fields (One per row)
    drawable.text((label_x, y1), "Family:", fill=0, font=font_bold)
    drawable.text(
        (label_x + value_offset, y1),
        printable_data.modelfamily,
        fill=0,
        font=font_regular,
    )

    drawable.text((label_x, y2), "Model:", fill=0, font=font_bold)
    drawable.text(
        (label_x + value_offset, y2),
        printable_data.modelname,
        fill=0,
        font=font_regular,
    )

    drawable.text((label_x, y3), "Hours:", fill=0, font=font_bold)
    drawable.text(
        (label_x + value_offset, y3),
        printable_data.power_on_hours,
        fill=0,
        font=font_regular,
    )

    # Cycles and Errors on the same line
    drawable.text((label_x, y4), "Cycles:", fill=0, font=font_bold)
    drawable.text(
        (label_x + value_offset, y4),
        printable_data.power_cycle,
        fill=0,
        font=font_regular,
    )

    drawable.text(
        (label_x + right_field_spacing, y4), "Errors:", fill=0, font=font_bold
    )
    drawable.text(
        (label_x + right_field_spacing + value_offset, y4),
        printable_data.smart_error_count,
        fill=0,
        font=font_regular,
    )

    # Continue remaining fields
    drawable.text((label_x, y5), "Shred on:", fill=0, font=font_bold)
    drawable.text(
        (label_x + value_offset, y5),
        printable_data.shred_timestamp,
        fill=0,
        font=font_regular,
    )

    drawable.text((label_x, y6), "Duration:", fill=0, font=font_bold)
    drawable.text(
        (label_x + value_offset, y6),
        printable_data.shred_duration,
        fill=0,
        font=font_regular,
    )

    # Capacity at the bottom
    drawable.text(
        (x_capacity, y_capacity),
        printable_data.capacity,
        fill=0,
        font=font_bold_bigger,
    )

    if (printable_data.connectiontype == "sata"):
        drawable.text(
            (x_connection_type, y_connection_type),
            "⬤ SATA",
            fill=0,
            font=font_regular,
        )

        drawable.text(
            (x_connection_type, y_connection_type+y_spacing_connection_type),
            "◯ NVME",
            fill=0,
            font=font_regular,
        )
    elif (printable_data.connectiontype == "nvme"):
        drawable.text(
            (x_connection_type, y_connection_type),
            "◯ SATA",
            fill=0,
            font=font_regular,
        )

        drawable.text(
            (x_connection_type, y_connection_type+y_spacing_connection_type),
            "⬤ NVME",
            fill=0,
            font=font_regular,
        )
    else:
        drawable.text(
            (x_connection_type, y_connection_type),
            "◯ SATA",
            fill=0,
            font=font_regular,
        )

        drawable.text(
            (x_connection_type, y_connection_type+y_spacing_connection_type),
            "◯ NVME",
            fill=0,
            font=font_regular,
        )


def draw_qr_code(image, data):
    """
    Draws a QR code on the provided image at a specific location.

    Parameters:
        image (Image.Image): The target image.
        data (str): The data to encode in the QR code.
    """
    # Generate the QR code
    qr_img = qrcode.make(data)
    qr_img = qr_img.convert("1")  # Ensure QR code is in binary (black/white)

    # Remove white border
    bbox = qr_img.getbbox()  # Get the bounding box of the non-white area
    qr_img = qr_img.crop(bbox)

    # Resize to desired size
    qr_img = qr_img.resize((QR_CODE_SIZE, QR_CODE_SIZE), Image.Resampling.LANCZOS)

    # Paste the QR code onto the image
    region = (5, 5, 5 + QR_CODE_SIZE, 5 + QR_CODE_SIZE)
    image.paste(qr_img, box=region)

def draw_outline(drawable, margin, width, output_width, output_height):
    """
    Draws a rectangular outline on the drawable canvas.

    Parameters:
        drawable (ImageDraw.Draw): The drawable canvas.
        margin (int): The margin from the edges of the image.
        width (int): The width of the outline.
        output_width (int): The total width of the image.
        output_height (int): The total height of the image.
    """
    # Define the four corners of the rectangle, adjusting for line width
    top_left = (margin, margin)
    top_right = (output_width - margin - width, margin)
    bottom_left = (margin, output_height - margin - width)
    bottom_right = (output_width - margin - width, output_height - margin - width)

    # Draw the outline lines with adjusted coordinates
    lines = [
        (top_left, top_right),  # Top edge
        (top_left, bottom_left),  # Left edge
        (top_right, bottom_right),  # Right edge
        (bottom_left, bottom_right),  # Bottom edge
    ]

    for line in lines:
        drawable.line(line, fill=0, width=width)

def generate_image(drive, rehdd_info, output_file):
    """Generates an image containing drive data and a QR code."""
    try:

        drive_json = DriveDataJson(
            state=drive.drive_state,
            contype=drive.drive_connection_type,
            fam=drive.modelfamily,
            name=drive.modelname,
            cap=drive.capacity,
            sn=drive.serialnumber,
            poh=drive.power_on_hours,
            pc=drive.power_cycle,
            err=drive.smart_error_count,
            time=int(drive.shred_timestamp),
            dur=drive.shred_duration,
        )

        qr_data = json.dumps(dataclasses.asdict(QrDataJson(drive_json, rehdd_info)))
        printable_data = format_to_printable(drive)
    except Exception as e:
        logging.error(f"Error preparing data: {e}")
        return

    output_image = Image.new("1", (OUTPUT_WIDTH, OUTPUT_HEIGHT), "white")
    draw = ImageDraw.Draw(output_image)

    font_regular = ImageFont.truetype(find_font_path(DEFAULT_FONT_REGULAR), 20)
    font_bold = ImageFont.truetype(find_font_path(DEFAULT_FONT_BOLD), 20)
    font_bold_bigger = ImageFont.truetype(find_font_path(DEFAULT_FONT_BOLD), 42)

    draw_outline(draw, 0, 3, OUTPUT_WIDTH + 1, OUTPUT_HEIGHT + 1)
    draw_text(draw, printable_data, font_regular, font_bold, font_bold_bigger)
    draw_qr_code(output_image, qr_data)

    try:
        output_image.save(output_file)
        logging.info(f"Image saved to {output_file}")
    except Exception as e:
        logging.error(f"Error saving image: {e}")


def main():
    rehdd_info = ReHddInfo("https://git.mosad.xyz/localhorst/reHDD", "bV0.2.2")

    temp_drive = DriveData(
        drive_index=0,
        drive_connection_type="sata",
        drive_state="shredded",
        modelfamily='Toshiba 2.5" HDD MK..65GSSX',
        modelname="TOSHIBA MK3265GSDX",
        capacity=343597383000,
        serialnumber="YG6742U56UDRL123456789ABCDEFGJKL",
        power_on_hours=7074,
        power_cycle=4792,
        smart_error_count=1,
        shred_timestamp=datetime.datetime.now(datetime.timezone.utc).timestamp(),
        shred_duration=81718,
    )

    generate_image(temp_drive, rehdd_info, "output.png")


if __name__ == "__main__":
    main()