mboot.py

import os
import subprocess
import sys

# Directories
SRC_DIR = os.path.abspath('./')
OBJ_DIR = os.path.abspath('./bin')
BOOT_DIR = os.path.abspath('./boot')
ISO_SCRIPT = os.path.abspath('./scripts/makeiso.sh')
RUN_SCRIPT = os.path.abspath('./scripts/run.sh')

# Tools
CC = 'gcc'
LD = 'ld'
ASM = 'nasm'

# Compiler flags
CFLAGS = ['-c', '-ffreestanding', '-m32', '-O0', '-Wall', '-Wextra', '-Werror', '-g', '-I', 'include/']  # Now a list of separate flags
ASMFLAGS = '-f elf32'
NOASMFLAGS = '-f bin'
LDFLAGS = '-T linker.ld -m elf_i386'

# Source files
ASM_SOURCES = []
C_SOURCES = []

# Object files for asm and c separately
ASM_OBJECTS = [os.path.join(OBJ_DIR, f.replace('.asm', '.asm.o')) for f in ASM_SOURCES]
C_OBJECTS = [os.path.join(OBJ_DIR, f.replace('.c', '.c.o')) for f in C_SOURCES]

# Target binary
TARGET = os.path.join(OBJ_DIR, 'loader.elf')
FINAL  = os.path.join(OBJ_DIR, 'loader.bin')


def move_element(lst:list[str], value:str, index:int):
    """Moves element x to position y in the list, if it exists."""
    try:
        # Find the index of element x
        index_x = lst.index(value)
        
        # Remove the element from its current position
        lst.pop(index_x)
        
        # Insert the element at the desired position y
        lst.insert(index, value)
    except ValueError:
        print(f"Element {value} not found in the list.")
    except IndexError:
        print(f"Invalid position {index}. List size: {len(lst)}")


def get_sources(walk_dir:str|None=None) -> tuple[list[str], list[str]]:
    """Recursively get all source files (.asm and .c) while ignoring directories starting with '.'."""
    if not walk_dir:
        walk_dir = SRC_DIR

    asm_files = []
    c_files = []
    for root, dirs, files in os.walk(walk_dir):
        # Modify dirs in-place to skip any directory that starts with a dot
        dirs[:] = [d for d in dirs if not d.startswith('.')]
        for file in files:
            if file.endswith('.asm'):
                asm_files.append(os.path.abspath(os.path.join(root, file)))
            elif file.endswith('.c'):
                c_files.append(os.path.abspath(os.path.join(root, file)))
    return asm_files, c_files


def run_command(command):
    """Helper function to run a shell command with modified environment"""
    tcmd = ''
    for cmd in command:
        try:
            tcmd = f"{tcmd} {cmd}"
        except:
            tcmd = f"{tcmd} {' '.join(cmd)}"
    print(f"Running: {tcmd}")
    env = os.environ.copy()  # Copy the current environment
    subprocess.run(tcmd, shell=True, env=env)

def build_asm():
    """Compile the ASM files"""
    for asm_file, obj_file in zip(ASM_SOURCES, ASM_OBJECTS):
        run_command([ASM, ASMFLAGS, os.path.join(SRC_DIR, asm_file), '-o', obj_file])

def build_boot_asm(boot_src, boot_out):
    """Compile the ASM files"""
    for asm_file, obj_file in zip(boot_src, boot_out):
        run_command([ASM, NOASMFLAGS, os.path.join(f'{SRC_DIR}.noasm/', asm_file), '-o', obj_file])

def build_c():
    """Compile the C files"""
    for c_file, obj_file in zip(C_SOURCES, C_OBJECTS):
        run_command([CC] + CFLAGS + ['-c', os.path.join(SRC_DIR, c_file), '-o', obj_file])

def link():
    """Link the object files into a final binary"""
    run_command([LD, LDFLAGS, '-o', TARGET] + ASM_OBJECTS + C_OBJECTS)
    run_command(['objcopy', '-O' 'binary', TARGET, FINAL])

def create_iso():
    """Create the ISO image"""
    run_command(["sh", ISO_SCRIPT])

def run_os():
    """Run the OS with the provided script"""
    run_command(["sh", RUN_SCRIPT])

def clean():
    """Clean up the build directory"""
    if not os.path.exists(OBJ_DIR):
        print(f"Directory '{OBJ_DIR}' does not exist.")
        return
    
    if os.path.exists(OBJ_DIR):
        # List all items (files and directories) in the directory
        for item in os.listdir(OBJ_DIR):
            item_path = os.path.join(OBJ_DIR, item)
            
            # If it's a directory, remove its contents recursively
            if os.path.isdir(item_path):
                for sub_item in os.listdir(item_path):
                    sub_item_path = os.path.join(item_path, sub_item)
                    if os.path.isdir(sub_item_path):
                        os.rmdir(sub_item_path)  # Remove empty subdirectory
                    else:
                        os.remove(sub_item_path)  # Remove file
                os.rmdir(item_path)  # Remove the empty directory
            else:
                os.remove(item_path)  # If it's a file, remove it

        print(f"Directory '{OBJ_DIR}' has been cleaned.")
    else:
        print(f"Directory '{OBJ_DIR}' does not exist.")
        

    """Clean up the second build directory"""
    if not os.path.exists(BOOT_DIR):
        print(f"Directory '{BOOT_DIR}' does not exist.")
        return
    
    if os.path.exists(BOOT_DIR):
        # List all items (files and directories) in the directory
        for item in os.listdir(BOOT_DIR):
            item_path = os.path.join(BOOT_DIR, item)
            
            # If it's a directory, remove its contents recursively
            if os.path.isdir(item_path):
                for sub_item in os.listdir(item_path):
                    sub_item_path = os.path.join(item_path, sub_item)
                    if os.path.isdir(sub_item_path):
                        os.rmdir(sub_item_path)  # Remove empty subdirectory
                    else:
                        os.remove(sub_item_path)  # Remove file
                os.rmdir(item_path)  # Remove the empty directory
            else:
                os.remove(item_path)  # If it's a file, remove it

        print(f"Directory '{BOOT_DIR}' has been cleaned.")
    else:
        print(f"Directory '{BOOT_DIR}' does not exist.")
        
    if os.path.exists(TARGET):
        os.remove(TARGET)


def main():
    global ASM_SOURCES, C_SOURCES, ASM_OBJECTS, C_OBJECTS
    if len(sys.argv) < 2:
        print("Usage: make.py [all|iso|run|clean]")
        sys.exit(1)

    command = sys.argv[1]
    
    # Get all source files (in subdirectories)
    asm_sources, c_sources = get_sources()

    # Create object file paths based on source files
    ASM_SOURCES = [_ for _ in asm_sources]
    C_SOURCES = [_ for _ in c_sources]
    ASMBOOT = get_sources('./.noasm/')[0]

    move_element(ASM_SOURCES, os.path.abspath("./src/asm/kernel_entry_loader.asm"), 0)
    move_element(C_SOURCES, os.path.abspath("./src/core/kernel_loader.c"), 0)
            
    ASM_OBJECTS = [os.path.join(OBJ_DIR,  f.replace('/', '_').replace('.asm', '.asm.o')) for f in ASM_SOURCES]
    C_OBJECTS   = [os.path.join(OBJ_DIR,  f.replace('/', '_').replace('.c', '.c.o'))     for f in C_SOURCES]
    ASMOBJ      = [os.path.join(BOOT_DIR, f.replace('/', '_').replace('.asm', '.bin'))   for f in ASMBOOT]

    with open("log.log", 'w') as log:
        log.write(str(ASM_SOURCES))
        log.write(str(ASM_OBJECTS))
        log.write(str(C_SOURCES))
        log.write(str(C_OBJECTS))

    if command == 'all':
        os.makedirs(OBJ_DIR, exist_ok=True)
        build_asm()
        build_c()
        link()

        build_boot_asm(ASMBOOT, ASMOBJ)
    elif command == 'iso':
        create_iso()
    elif command == 'run':
        run_os()
    elif command == 'clean':
        clean()
    elif command == 'help' or command == '-help' or command == '--help' \
    or command == 'h' or command == '-h' or command == '--h':
        print("Usage: make.py [all|iso|run|clean]")
    else:
        print("Unknown command:", command)
        sys.exit(1)


if __name__ == '__main__':
    main()