diff --git a/roles/driver-msi-keyboard-color/files/keyboard_color.py b/roles/driver-msi-keyboard-color/files/keyboard_color.py
index 17480c6d..d38d4e4d 100644
--- a/roles/driver-msi-keyboard-color/files/keyboard_color.py
+++ b/roles/driver-msi-keyboard-color/files/keyboard_color.py
@@ -3,11 +3,29 @@ import datetime
 import subprocess
 
 def hex_to_rgb(hex_color):
-    """Converts Hex color to an RGB tuple."""
+    """ Converts a hexadecimal color string to an RGB tuple.
+    
+    Args:
+        hex_color (str): Hexadecimal color string.
+
+    Returns:
+        tuple: Tuple representing the RGB color.
+    """
+    if len(hex_color) != 6:
+        raise ValueError("Hex color must be 6 characters long.")
     return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
 
 def calculate_color(start_color, end_color, ratio):
-    """Calculates an interpolated color between two colors based on a transition ratio."""
+    """ Calculates the interpolated color between two colors.
+
+    Args:
+        start_color (str): Start color in hex format.
+        end_color (str): End color in hex format.
+        ratio (float): Interpolation ratio (0 to 1).
+
+    Returns:
+        str: Interpolated color in hex format.
+    """
     start_rgb = hex_to_rgb(start_color)
     end_rgb = hex_to_rgb(end_color)
     current_rgb = [round(start + (end - start) * ratio) for start, end in zip(start_rgb, end_rgb)]
@@ -25,21 +43,11 @@ def get_current_period(current_time, color_times):
     Returns:
     tuple: A tuple of the start and end color (as hex codes) for the current period.
     """
-    # Sort the list of time periods
     sorted_periods = sorted(color_times.items())
-
-    # Iterate over the sorted periods to find the current period
     for i, (period_start_str, colors) in enumerate(sorted_periods):
         period_start_time = datetime.datetime.strptime(period_start_str, "%H:%M").time()
-
-        # If the current time is before the start of the next period
-        if current_time < period_start_time:
-            # If we are in the first period, the previous period is the last of the day
-            previous_period_index = i - 1 if i > 0 else -1
-            return sorted_periods[previous_period_index][1]
-
-    # If the current time is after the last defined period,
-    # use the colors of the first period as default
+        if current_time <= period_start_time:
+            return sorted_periods[i - 1 if i > 0 else -1][1]
     return sorted_periods[0][1]
 
 def calculate_transition_ratio(current_time, start_time, end_time):
@@ -71,6 +79,20 @@ def calculate_transition_ratio(current_time, start_time, end_time):
     # Ratio as the proportion of time elapsed in the total transition duration
     return time_since_start / transition_duration
 
+def change_keyboard_color(color, device_id):
+    """ Changes the keyboard color using an external command.
+
+    Args:
+        color (str): Color in hex format.
+        device_id (str): Vendor and product ID for the keyboard.
+    """
+    command = ["msi-perkeyrgb", "--model", "GS65", "-s", color, "--id", device_id]
+    try:
+        subprocess.run(command, check=True)
+        print(f"Keyboard color changed to #{color}.")
+    except subprocess.CalledProcessError as e:
+        print(f"Error setting keyboard color: {e}")
+        sys.exit(1)
 
 def main(vendor_and_product_id):
     color_times = {
@@ -86,21 +108,26 @@ def main(vendor_and_product_id):
         "04:00": ("990000", "5bc0eb"),  # Night to Blue
     }
 
+    # Ensure 24-hour coverage
+    if "00:00" not in color_times:
+        raise ValueError("Time periods must cover the full 24-hour cycle.")
+
     current_time = datetime.datetime.now().time()
     start_color, end_color = get_current_period(current_time, color_times)
 
-    next_time = list(color_times.keys())[0]
-    next_start_time_obj = datetime.datetime.strptime(next_time, "%H:%M").time()
+    sorted_times = sorted(color_times.keys())
+    for time_str in sorted_times + [sorted_times[0]]:
+        next_start_time_obj = datetime.datetime.strptime(time_str, "%H:%M").time()
+        if current_time < next_start_time_obj:
+            break
 
-    transition_ratio = calculate_transition_ratio(current_time, start_color, next_start_time_obj)
+    current_period_start_time_str = [time for time, colors in color_times.items() if colors == (start_color, end_color)][0]
+    current_period_start_time_obj = datetime.datetime.strptime(current_period_start_time_str, "%H:%M").time()
+
+    transition_ratio = calculate_transition_ratio(current_time, current_period_start_time_obj, next_start_time_obj)
     current_color = calculate_color(start_color, end_color, transition_ratio)
 
-    print(f"Changing keyboard color to #{current_color}.")
-    try:
-        subprocess.run(["msi-perkeyrgb", "--model", "GS65", "-s", current_color, "--id", vendor_and_product_id], check=True)
-    except subprocess.CalledProcessError as e:
-        print(f"Error setting keyboard color: {e}")
-        sys.exit(1)
+    change_keyboard_color(current_color, vendor_and_product_id)
 
 if __name__ == "__main__":
     if len(sys.argv) < 2: