ArchInstaller

#!/usr/bin/env python
import concurrent.futures, yaml, requests, time, sys, os
from subprocess import run, PIPE

class TUI:
    OKPURPLE = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    ANIMATION= [
    "[        ]",
    "[=       ]",
    "[==      ]",
    "[===     ]",
    "[====    ]",
    "[=====   ]",
    "[======  ]",
    "[======= ]",
    "[========]",
    "[ =======]",
    "[  ======]",
    "[   =====]",
    "[    ====]",
    "[     ===]",
    "[      ==]",
    "[       =]",
    "[        ]",
    "[        ]"
    ]




class ArchInstaller:
    def __init__(self) -> None:
        try:
            if len(sys.argv) != 1:
                with open("./config.yaml", mode="w") as f:
                    response = requests.get(sys.argv[1])
                    f.write(response.text)
            with open("./config.yaml") as data:
                config = yaml.load(data, Loader=yaml.FullLoader)
                if os.environ.get("USER") == "root":
                    self.config = config["root"]
                else:
                    self.config = config["user"]
        except:
            print(f"{TUI.OKGREEN}==>{TUI.ENDC} There is a problem with the config.yaml")
            exit()
        try:
            requests.get("https://google.com/")
        except:
            print(f"{TUI.OKGREEN}==>{TUI.ENDC} Your connection is not stable")
            exit()

    def ascii(self) -> None:
        ASCII = """    ___    ____  ________  _______   ________________    __    __    __________ 
   /   |  / __ \/ ____/ / / /  _/ | / / ___/_  __/   |  / /   / /   / ____/ __ \\
  / /| | / /_/ / /   / /_/ // //  |/ /\__ \ / / / /| | / /   / /   / __/ / /_/ /
 / ___ |/ _, _/ /___/ __  // // /|  /___/ // / / ___ |/ /___/ /___/ /___/ _, _/ 
/_/  |_/_/ |_|\____/_/ /_/___/_/ |_//____//_/ /_/  |_/_____/_____/_____/_/ |_|  
"""
        print(TUI.OKGREEN + ASCII + TUI.ENDC)

    def exec(self,command: str, description: str|None = None , chroot:bool = False):
        if chroot:
            sh = 'arch-chroot /mnt ' + command  
        else:
            sh = command
        with concurrent.futures.ThreadPoolExecutor() as executor:
            future = executor.submit(lambda: run(sh, stdout=PIPE, stderr=PIPE, shell=True, encoding='utf-8'))
            if description != None:
                i = 0
                while future.running():
                    print(description + ' ' + TUI.OKPURPLE + TUI.ANIMATION[i % len(TUI.ANIMATION)] + TUI.ENDC, end='\r')
                    time.sleep(.1)
                    i += 1
                if future.result().returncode == 0:
                    print(description + f' {TUI.OKGREEN}✔{TUI.ENDC}' + "                            ", end='\n')
                else:
                    print(description + f' {TUI.FAIL}✘{TUI.ENDC}' + "                               ", end='\n')
    
    def setNtp(self) -> None:
        description = f"{TUI.OKGREEN}==>{TUI.ENDC} set ntp to {self.config['settings']['ntp']}"
        command = f"timedatectl set-ntp {self.config['settings']['ntp']}"
        self.exec(command,description)

    def driveConfigurator(self) -> None:
        print(f"{TUI.OKGREEN}==>{TUI.ENDC} Config Hard Drive")  # with fdisk

        drive = self.config["settings"]["drive"]

        # Format Hard Drive
        if drive["erase"]:
            description = f"    {TUI.OKBLUE}-->{TUI.ENDC} Formatting Drive : {drive['blk']}"
            command = f"echo 'g\nw\n' | fdisk {drive['blk']}"
            self.exec(command, description)

        # Making Partitions
        print(f"    {TUI.OKBLUE}-->{TUI.ENDC} Create Partitions")
        for partition in drive["partitions"]:
            size = partition["size"]
            number = partition["number"]
            label = partition["label"]
            filesystem = partition["filesystem"]
            description = f"        {TUI.OKBLUE}-->{TUI.ENDC} Partition : {partition['label']}"
            if partition["number"] == 1:
                command = f"echo 'n\n{number}\n\n{size}\nt\n{filesystem}\nw\n' | fdisk {drive['blk']}"
            else:
                command = f"echo 'n\n{number}\n\n{size}\nt\n{number}\n{filesystem}\nw\n' | fdisk {drive['blk']}"
            self.exec(command, description)
            
            # Make FileSystem & add Label to Partition
            description = f"            {TUI.OKPURPLE}-->{TUI.ENDC} Make filesystem : {filesystem}"
            match filesystem:
                case "uefi":
                    self.exec(f"mkfs.fat -F32 {drive['blk']}{number}")
                    command = f"fatlabel {drive['blk']}{number} {label}"
                case "linux":
                    command = f"echo 'y\n' | mkfs.ext4 -L {label} {drive['blk']}{number}"
                case "swap":
                    command = f"mkswap {drive['blk']}{number}"
            self.exec(command,description)

    def mount(self):

        drive = self.config["settings"]["drive"]
        # Mount Partitions
        print(f"{TUI.OKGREEN}==>{TUI.ENDC} Mount Partitions ")

        # root -> /mnt
        while not os.path.exists("/dev/disk/by-label/ROOT"):
            time.sleep(1)
        description = f"    {TUI.OKBLUE}-->{TUI.ENDC} ROOT"
        command = f"mount /dev/disk/by-label/ROOT /mnt"
        self.exec(command,description)

        self.exec("mkdir -p /mnt/{home,boot/efi}")

        for partition in drive["partitions"]:
            
            # boot -> /mnt/boot/efi
            if partition["label"] == "BOOT":
                description = f"    {TUI.OKBLUE}-->{TUI.ENDC} BOOT"
                command = f"mount {drive['blk']}{partition['number']} /mnt/boot/efi"
                self.exec(command,description)    
            
            # Mount Home Partition if exist
            # home -> /mnt/home
            if partition["label"] == "HOME":
                description = f"    {TUI.OKBLUE}-->{TUI.ENDC} HOME"
                command = f"mount {drive['blk']}{partition['number']} /mnt/home"
                self.exec(command,description)

            # enable swap if exists
            if partition["label"] == "SWAP":
                description = f"    {TUI.OKBLUE}-->{TUI.ENDC} Enable SWAP Partition "
                command = f"swapon {drive['blk']}{partition['number']}"
                self.exec(command,description)

    def installPackages(self) -> None:
        print(f"{TUI.OKGREEN}==>{TUI.ENDC} Install Packages")
        if os.environ.get('USER') == "root":
            for package in self.config["packages"]:
                description = f"    {TUI.OKBLUE}-->{TUI.ENDC} {package}"
                command = f"pacstrap /mnt {package}"
                self.exec(command,description)
        else:
            for package in self.config["packages"]:
                description = f"    {TUI.OKBLUE}-->{TUI.ENDC} {package}"
                command = f"{self.config['aurhelper']} -S --skipreview --noconfirm {package}"
                self.exec(command,description)

    def genFstab(self) -> None:
        print(f"{TUI.OKGREEN}==>{TUI.ENDC} Generating FSTAB {TUI.OKGREEN}✔{TUI.ENDC}")
        with open('/mnt/etc/fstab', 'w') as fstab:
            run(['genfstab', '-U', '/mnt'], stdout=fstab)

    def setLocale(self) -> None:
        print(f"{TUI.OKGREEN}==>{TUI.ENDC} Setting Locale")
        with open("/mnt/etc/locale.gen", "a") as f:
            for locale in self.config["settings"]["locale"]:
                print(f"    {TUI.OKBLUE}-->{TUI.ENDC} Generating {locale} {TUI.OKGREEN}✔{TUI.ENDC}")
                f.write(f"{locale}\n")
        self.exec("locale-gen", chroot=True)
        print(f"    {TUI.OKBLUE}-->{TUI.ENDC} set LANG env to {self.config['settings']['lang']} {TUI.OKGREEN}✔{TUI.ENDC}")
        with open("/mnt/etc/locale.conf", "w") as f:
            f.write(f"LANG={self.config['settings']['lang']}\n")

    def setTimeZone(self) -> None:
        description = f"{TUI.OKGREEN}==>{TUI.ENDC} setting TimeZone to {self.config['settings']['timezone']}"
        command = f"ln -fs /usr/share/zoneinfo/{self.config['settings']['timezone']} /etc/localtime"
        self.exec(command, description, True)
        self.exec("hwclock --systohc --utc", chroot=True)

    def installBootLoader(self):
        description = f"{TUI.OKGREEN}==>{TUI.ENDC} Generating BootLoader"
        command = f"grub-install --recheck"
        self.exec(command,description,True)
        self.exec(f"grub-mkconfig -o /boot/grub/grub.cfg", chroot=True)

    def setHostName(self) -> None:
        print(f"{TUI.OKGREEN}==>{TUI.ENDC} set hostname to {self.config['settings']['hostname']} {TUI.OKGREEN}✔{TUI.ENDC}")
        with open("/mnt/etc/hostname", "w") as f:
            f.write(self.config["settings"]["hostname"])

    def setHosts(self) -> None:
        print(f"{TUI.OKGREEN}==>{TUI.ENDC} set hosts {TUI.OKGREEN}✔{TUI.ENDC}")
        with open("/mnt/etc/hosts", "a") as f:
            f.write(
                f"127.0.0.1     localhost.localdomain        localhost\n::1           localhost.localdomain        localhost\n127.0.1.1    {self.config['settings']['hostname']}.localdomain    {self.config['settings']['hostname']}"
            )

    def enableServices(self):
        print(f"{TUI.OKGREEN}==>{TUI.ENDC} Enabling Services")
        for service in self.config["services"]:
            description = f"    {TUI.OKBLUE}-->{TUI.ENDC} {service}"
            if os.environ.get('USER') == "root":
                command = f"systemctl enable {service}"
                chroot = True
            else:
                command = f"sudo systemctl enable {service}"
                chroot = False
            self.exec(command,description,chroot)

    def setRootPassword(self):
        description = f"{TUI.OKGREEN}==>{TUI.ENDC} Set Root Password"
        command = f"echo -n 'root:password' | chpasswd"
        self.exec(command, description, True)

    def createUser(self):
        username = self.config["username"]
        shell = self.config["settings"]["shell"]
        description = f"{TUI.OKGREEN}==>{TUI.ENDC} Creating User ({username})"
        command = f"useradd -m -g users -G wheel -s {shell} {username}"
        self.exec(command, description, True)
        self.exec(f"echo '{username}:password' | chpasswd", chroot=True)

    def sudoersConf(self, conf):
        if conf == "UA":
            self.exec("chmod 644 /etc/sudoers", chroot=True)
            description = f"{TUI.OKGREEN}==>{TUI.ENDC} Config sudoers file"
            command = r"sed -i 's/^#\s*\(%wheel\s*ALL=(ALL:ALL)\s*ALL\)/\1/' /mnt/etc/sudoers"
            self.exec(command,description)
        elif conf == "UN":
            self.exec(r"sudo sed -i 's/^#\s*\(%wheel\s*ALL=(ALL:ALL)\s*NOPASSWD:\s*ALL\)/\1/' /etc/sudoers")
        elif conf == "CN":
            self.exec(r"sudo sed -i 's/^\s*\(%wheel\s*ALL=(ALL:ALL)\s*NOPASSWD:\s*ALL\)/# \1/' /etc/sudoers")

    def copyInstaller(self):
        self.exec(f"mkdir -p /mnt/home/{self.config['username']}/archinstaller")
        self.exec(f"cp -r . /mnt/home/{self.config['username']}/archinstaller")
        self.exec(f"chown -R {self.config['username']}:wheel /mnt/home/{self.config['username']}/archinstaller", chroot=True)
    
    def installDependencies(self):
        description = f"{TUI.OKGREEN}==>{TUI.ENDC} Install dependencies"
        depends = ['git', 'python-yaml', 'python-requests']
        command = "pacstrap /mnt " + " ".join(depends)
        self.exec(command, description)

    def umount(self):
        description = f"{TUI.OKGREEN}==>{TUI.ENDC} Umount Partitions"
        command = "umount -R /mnt"
        self.exec(command, description)

    def reboot(self):
        print("====================================\nInstallation completed!!")
        input("Press enter to reboot")
        self.exec("reboot")

    def installAURHelper(self):
        self.exec(f"cd /tmp && git clone http://aur.archlinux.org/{self.config['aurhelper']}-bin.git")
        description = f"{TUI.OKGREEN}==>{TUI.ENDC} Install AUR helper"
        command = f"cd /tmp/{self.config['aurhelper']}-bin && makepkg -si --noconfirm"
        self.exec(command,description)

    def runUserScripts(self):
        print(f"{TUI.OKGREEN}==>{TUI.ENDC} Execute Your Scripts")
        for script in self.config["scripts"]:
            description = f"    {TUI.OKBLUE}-->{TUI.ENDC} {script['name']}"
            if 'shell' in script:
                command = script['shell']
            elif 'script' in script:
                self.exec(f"chmod +x {script['script']}")
                command = script["script"]
            else:
                description += " : Yaml file is incorret"
                command = "fakesudo :)"
            self.exec(command, description)

    def run(self):
        self.ascii()
        if os.environ.get('USER') == "root":
            self.setNtp()
            self.driveConfigurator()
            self.mount()
            self.installPackages()
            self.genFstab()
            self.setLocale()
            self.setTimeZone()
            self.installBootLoader()
            self.setHostName()
            self.setHosts()
            self.enableServices()
            self.setRootPassword()
            self.createUser()
            self.sudoersConf("UA")
            self.copyInstaller()
            self.installDependencies()
            self.umount()
            self.reboot()
        else:
            self.sudoersConf("UN")
            self.installAURHelper()
            self.installPackages()
            self.runUserScripts()
            self.enableServices()
            self.sudoersConf("CN")

ArchInstaller().run()