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 = 300  # px
TEXT_X_OFFSET = 300  # px
QR_CODE_SIZE = 289  # 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
    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 DriveDataPrintable:
    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 DriveDataJson:
    drive: DriveData
    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(
        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(16, 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 the formatted text onto the image."""
    text_y_offset = 10
    value_column_x_offset = 120
    font_size = 20

    drawable.text(
        (TEXT_X_OFFSET, text_y_offset), printable_data.serialnumber, (0), font=font_bold
    )
    text_y_offset += 25

    fields = [
        ("Family:", printable_data.modelfamily, font_regular),
        ("Model:", printable_data.modelname, font_regular),
        ("Hours:", printable_data.power_on_hours, font_regular),
        ("Cycles:", printable_data.power_cycle, font_regular),
        ("Errors:", printable_data.smart_error_count, font_regular),
        ("Shred on:", printable_data.shred_timestamp, font_regular),
        ("Duration:", printable_data.shred_duration, font_regular),
    ]

    for label, value, font in fields:
        drawable.text((TEXT_X_OFFSET, text_y_offset), label, fill=0, font=font_bold)
        drawable.text(
            (TEXT_X_OFFSET + value_column_x_offset, text_y_offset),
            value,
            fill=0,
            font=font,
        )
        text_y_offset += 25

    drawable.text(
        (TEXT_X_OFFSET, text_y_offset),
        printable_data.capacity,
        (0),
        font=font_bold_bigger,
    )


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:
        qr_data = json.dumps(dataclasses.asdict(DriveDataJson(drive, 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), 60)

    draw_outline(draw, 1, 4, OUTPUT_WIDTH, OUTPUT_HEIGHT)
    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_state="shredded",
        modelfamily='Toshiba 2.5" HDD MK..65GSSX',
        modelname="TOSHIBA MK3265GSDX",
        capacity=343597383680,
        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()