Fix /etc/apt/sources.list.d/ sources files with suite/release and key updates.

apt_fix_sources.py

#!/usr/bin/env python3
# Fix /etc/apt/sources.list.d/ sources files with suite/release and key updates.

# Author: Maulik Mistry with AI support.
# Please share support: https://www.paypal.com/paypalme/m1st0

# License: BSD License 2.0
# Copyright (c) 2025, Maulik Mistry
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the <organization> nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import sys
import shutil
import re
import requests
import difflib
import argparse
from pathlib import Path
from rich.console import Console
from rich.markup import escape

console = Console()

def fetch_key(key_source):
    try:
        if key_source.startswith(('http://', 'https://')):
            response = requests.get(key_source)
            response.raise_for_status()
            return response.text.splitlines()
        else:
            # Handle file path fetching
            key_path = Path(key_source).expanduser().resolve()
            if not key_path.is_file():
                raise RuntimeError(f"The file at {key_path} does not exist or is not a valid file.")
            
            with open(key_path, 'r') as f:
                return f.read().splitlines()  # Read file contents and split into lines
    except Exception as e:
        raise RuntimeError(f"Failed to fetch key from {key_source}: {e}")        

def format_key(key_lines):
    if not key_lines:
        return []

    formatted = []
    # First line should always have "Signed-By:"
    formatted.append(f"Signed-By: {key_lines[0].strip()}")

    # Process the rest of the key lines
    for line in key_lines[1:]:
        stripped = line.strip()
        if stripped == "":
            # Only insert a dot if the line is completely empty after stripping
            formatted.append("  .")
        else:
            # Otherwise, just add the line with the required indentation
            formatted.append(f"  {stripped}")
    
    return formatted

def generate_new_content(original_lines, key_lines):
    output_lines = []
    skip_signed_by = False

    for line in original_lines:
        stripped = line.lstrip()
        if stripped.startswith('Enabled: no'):
            continue
        if skip_signed_by:
            if not line.startswith((' ', '\t')):
                skip_signed_by = False
            else:
                continue
        if re.match(r'^\s*Signed-By:', line):
            skip_signed_by = True
            continue
        output_lines.append(line.rstrip('\n'))

    insert_idx = None
    for idx, line in enumerate(output_lines):
        if line.strip().startswith('Components:'):
            insert_idx = idx + 1
            break
    if insert_idx is None:
        insert_idx = len(output_lines)

    new_lines = output_lines[:insert_idx] + format_key(key_lines) + output_lines[insert_idx:]
    return [line + '\n' for line in new_lines]

def show_diff(original_lines, new_lines, file_path):
    diff = difflib.unified_diff(
        original_lines,
        new_lines,
        fromfile=f"{file_path} (original)",
        tofile=f"{file_path} (new)",
        lineterm=""
    )
    for line in diff:
        if line.startswith('+') and not line.startswith('+++'):
            console.print(f"[green]{escape(line)}[/green]")
        elif line.startswith('-') and not line.startswith('---'):
            console.print(f"[red]{escape(line)}[/red]")
        elif line.startswith('@@'):
            console.print(f"[cyan]{escape(line)}[/cyan]")
        else:
            console.print(escape(line))

def backup_and_save(file_path, new_lines):
    backup_path = file_path.with_suffix(file_path.suffix + '.bak')
    shutil.copy2(file_path, backup_path)
    console.print(f"[green]Backup created at {backup_path}[/green]")
    with open(file_path, 'w') as f:
        f.writelines(new_lines)
    console.print(f"[green]Updated {file_path}[/green]\n")

def prompt_input(prompt, default=None):
    if default:
        prompt = f"{prompt} [{default}] "
    else:
        prompt = f"{prompt} "
    value = input(prompt).strip()
    return value if value else default

def process_file(file_path, retry_fetch=True):
    with open(file_path, 'r') as f:
        original_lines = f.readlines()

    # Show URI to user
    uri = extract_uri(original_lines)
    if uri:
        console.print(f"\n[bold]Found URI:[/bold] {uri}")

    # Check Suites/Suite
    suite_key, suite_value = extract_suite(original_lines)
    if suite_value:
        console.print(f"[bold]Found {suite_key}:[/bold] {suite_value}")
        if suite_value.lower() == "oracular":
            fix = prompt_input(f"Change {suite_key} from 'oracular' to 'plucky'? [Y/n]", "Y").lower()
            if fix == 'y':
                # Replace the suite line
                for idx, line in enumerate(original_lines):
                    if line.strip().startswith((suite_key + ":")):
                        original_lines[idx] = f"{suite_key}: plucky\n"
                        break
        else:
            fix = prompt_input(f"Suite '{suite_value}' found. Change to 'plucky'? [y/N]", "N").lower()
            if fix == 'y':
                for idx, line in enumerate(original_lines):
                    if line.strip().startswith((suite_key + ":")):
                        original_lines[idx] = f"{suite_key}: plucky\n"
                        break

    while True:
        key_source = input(f"\nEnter key URL or file path for [bold]{file_path.name}[/bold] (ENTER to skip): ").strip()
        if not key_source:
            console.print("[yellow]Skipping file.[/yellow]\n")
            return 'skipped'

        try:
            key_lines = fetch_key(key_source)
            break
        except Exception as e:
            console.print(f"[red]{e}[/red]")
            if not retry_fetch or prompt_input("Retry fetching key? [Y/n]", "Y").lower() != 'y':
                console.print("[yellow]Skipping file.[/yellow]\n")
                return 'skipped'

    new_lines = generate_new_content(original_lines, key_lines)

    console.print("\n[cyan]--- Diff Preview ---[/cyan]")
    show_diff(original_lines, new_lines, file_path)
    console.print("[cyan]--------------------[/cyan]\n")

    confirm = prompt_input("Apply changes? [Y/n]", "Y").lower()
    if confirm == 'y':
        backup_and_save(file_path, new_lines)
        return 'modified'
    else:
        console.print("[yellow]Changes discarded.[/yellow]\n")
        return 'skipped'

def extract_uri(lines):
    for line in lines:
        if line.strip().startswith("URIs:"):
            return line.strip().split(":", 1)[1].strip()
    return None

def extract_suite(lines):
    for line in lines:
        if line.strip().startswith(("Suite:", "Suites:")):
            key, value = line.strip().split(":", 1)
            return key, value.strip()
    return None, None

def main():
    parser = argparse.ArgumentParser(
        description="Fix .sources files by removing old Signed-By/Enabled lines and inserting new keys."
    )
    parser.add_argument(
        "directory",
        type=Path,
        help="Directory containing .sources files (e.g., /etc/apt/sources.list.d/)"
    )

    args = parser.parse_args()
    sources_dir = args.directory.resolve()

    if not sources_dir.is_dir():
        console.print(f"[red]{sources_dir} is not a valid directory![/red]")
        sys.exit(1)

    sources_files = sorted(sources_dir.glob('*.sources'))
    if not sources_files:
        console.print("[yellow]No .sources files found.[/yellow]")
        sys.exit(0)

    total_files = len(sources_files)
    modified = 0
    skipped = 0

    console.print(f"\nFound {total_files} .sources file(s) in [bold]{sources_dir}[/bold]\n")

    for idx, file_path in enumerate(sources_files, start=1):
        console.rule(f"File {idx}/{total_files}: {file_path.name}")
        choice = prompt_input("Process this file? [Y/n/q to quit]", "Y").lower()
        if choice == 'y':
            result = process_file(file_path)
            if result == 'modified':
                modified += 1
            else:
                skipped += 1
        elif choice == 'q':
            console.print("[bold]Exiting early by user request.[/bold]")
            break
        else:
            console.print("[yellow]Skipping.[/yellow]\n")
            skipped += 1

    # 🎉 Final Summary
    console.rule("[bold cyan]Summary[/bold cyan]")
    console.print(f"[green]Modified:[/green] {modified}")
    console.print(f"[yellow]Skipped:[/yellow] {skipped}")
    console.print(f"[blue]Total processed:[/blue] {modified + skipped}\n")

if __name__ == "__main__":
    main()