#!/bin/bash

function header () {
    tput setaf 1
    tput bold
    echo $@
    tput sgr0
    return
}
function help() {
    echo Usage: ${0} '[OPTIONS]'
    echo '\-d DISK -- Use the disk.'
    echo '\-e      -- Encrypt the root partition'
    echo '\-g      -- GUI packages and setup'
    echo '\-h      -- This helptext'
    echo '\-k      -- Kali Linux-like package additions'
    echo '\-p      -- Productivity package additions'
    echo '\-s      -- Create a layout for an AniNIX::Spartacus'
    echo '\-m      -- Skip disk operations and assume storage is mounted on /mnt'
    echo '\-z      -- Try to add all the packages on AniNIX::Core'
    exit 1;
}

spartacus=0;
encrypt=0;
gui=0;
kali=0;
kitchensink=0;
productivity=0;
disk="/dev/sda"
nodiskbuild=0;
bootpart=1;
rootpart=2;
datapart=99;
# TODO Add LVM as an argument
while getopts "d:egkpmsz" OPTION
do
    case $OPTION in
        d) disk=${OPTARG} ;;
        e) encrypt=1 ;;
        g) gui=1 ;;
        k) kali=1 ;;
        p) productivity=1; gui=1 ;;
        m) nodiskbuild=1 ;;
        s) spartacus=1 ;;
        z) kitchensink=1 ;;
        *) help
    esac
done

header Confirm options:
echo Spartacus set to: $spartacus
echo Encryption set to: $encrypt
echo GUI: $gui
echo Productivity: $productivity
echo Kali tools: $kali
echo All Core packages: $kitchensink
echo Disk to use: $disk \(Skip disk building? $nodiskbuild \)
printf "Is this OK? Type YES to continue: "
read answer
if [ "$answer" != "YES" ]; then
    echo User did not confirm.
    exit 1;
fi
pacman -Syy
if [ "$nodiskbuild" -eq 0 ]; then
    header Allocating space
    dd if=/dev/zero of="$disk" bs=1 count=2000000 # "$(fdisk -l "$disk" | head -n 1 | cut -f 5 -d ' ')"
    if [ "$spartacus" -eq 1 ]; then
        # Insert an ExFAT data partition ahead of the rest.
        export datapart=1;
        export bootpart=$((bootpart+1))
        export rootpart=$(($rootpart+1))
        # Break the disk up into 4ths -- 2/4 go to data, 1/4 go to boot, and 1/4 to root
        export disksize=$(($(fdisk -l $disk | head -n 1 | cut -f 5 -d ' ') / 1048576)) # Return disk size in MB
        if [ "$disksize" == "" ]; then echo "Can't identify disk size"; exit 1; fi
        if [ "$disksize" -lt 7788 ]; then echo "This drive is too small to be a Spartacus."; exit 1; fi # Must be 8GB or more to have 2GB root.
        export bootsize=$(($disksize / 4))
        export datasize=$(($disksize / 2))
        printf 'mklabel msdos\nmkpart primary ext4 1MiB %s\nmkpart primary ext4 %s %s\nmkpart primary ext4 %s 100%%FREE\nprint\nquit\n' $datasize"MiB" $datasize"MiB" $(($datasize+$bootsize))"MiB" $(($datasize+$bootsize))"MiB" | parted "$disk"
        #create data partition
        pacman -S exfat-utils --noconfirm
        mkfs.exfat "$disk""$datapart"
        exfatlabel "$disk""$datapart" "AS-XPLATFRM"
    else
        # One 200MB boot and the rest is root 
        printf 'mklabel msdos\nmkpart primary ext4 1MiB 201MiB\nmkpart primary ext4 513MiB 100%%FREE\nprint\nquit\n' | parted "$disk"
    fi 
    header Making ext4 boot partition on "$disk""$bootpart"
    mkfs.ext4 "$disk""$bootpart"
    tune2fs -L "BOOT" "$disk""$bootpart"

    header Making root and mountpoints
    if [ "$encrypt" -eq 1 ]; then 
        header Making encrypted root on "$disk""$rootpart"
        modprobe dm-crypt
        modprobe serpent_generic
        header Formatting root -- make sure to enter YES followed by a strong passphrase.
        cryptsetup luksFormat -c serpent-xts-plain64 -h sha512 --key-size 512 "$disk""$rootpart"
        header Unlocking root
        cryptsetup luksOpen "$disk""$rootpart" cryptroot
        mkfs.xfs -f /dev/mapper/cryptroot
        xfs_admin -L ROOT /dev/mapper/cryptroot
        mount /dev/mapper/cryptroot /mnt
        if [ $? -ne 0 ]; then header ERROR: Cannot continue; exit 1; fi
    else 
        header Making root on "$disk""$rootpart"
        mkfs.xfs -f "$disk""$rootpart"
        xfs_admin -L ROOT "$disk""$rootpart"
        mount "$disk""$rootpart" /mnt
        if [ $? -ne 0 ]; then header ERROR: Cannot continue; exit 1; fi
    fi

    mkdir /mnt/boot
    mount "$disk""$bootpart" /mnt/boot
    if [ "$?" -ne 0 ]; then header ERROR: Cannot continue; exit 1; fi
fi


# Install ArchLinux with basic clients for the AniNIX Services.
# * git for Foundation
# * lynx for WebServer and Wiki
# * openssh for SSH/SFTP
# * irssi for IRC
# * make for source packages
# * tor for anonymity
header Installing ArchLinux to device\(s\) on /mnt
export pkglist="base base-devel parted net-tools bind-tools git openssh make lynx irssi vim wget tor torsocks grub os-prober rsync openntpd tmux"
if [ "$gui" -eq 1 ]; then 
    export pkglist="$pkglist"" xorg-server xfce4 seamonkey conky tigervnc"
fi
if [ "$spartacus" -eq 1 ]; then
    export pkglist="$pkglist"" exfat-utils"
fi
if [ "$productivity" -eq 1 ]; then
    export pkglist="$pkglist"" libreoffice-still gimp feh vlc evince"
fi
if [ "$kali" -eq 1 ]; then
    export pkglist="$pkglist"" extundelete testdisk nmap tcpdump hexedit dcfldd"
    if [ "$gui" -eq 1 ]; then
        export pkglist="$pkglist"" wireshark-gtk"
    else 
        export pkglist="$pkglist"" wireshark-cli"
    fi
fi
if [ "$kitchensink" -eq 1 ]; then
    export pkglist="base base-devel $(wget -q -O - 'https://aninix.net/installed-packages.txt' | cut -f 1 -d ' ' | tr '\n' ' ')"
fi

yes "" | pacstrap -i /mnt $pkglist 
if [ $? -ne 0 ]; then header ERROR: Cannot continue -- pacstrap failed; exit 1; fi

header Create FSTAB
genfstab -U /mnt >> /mnt/etc/fstab

header Set time
sed -i 's/#en_US.UTF-8/en_US.UTF-8/' /mnt/etc/locale.gen
arch-chroot /mnt locale-gen
ln -sf /usr/share/zoneinfo/America/Chicago /mnt/etc/localtime
arch-chroot /mnt hwclock --systohc --utc

header Setup bootloader
if [ "$nodiskbuild" -eq 0 ]; then
    export rootuuid="$(blkid "$disk""$rootpart" | cut -f 2 -d '"')"
    if [ "$encrypt" -eq 1 ]; then
        export hookstring="$(grep 'HOOKS=' /mnt/etc/mkinitcpio.conf | grep -v '#')"
        sed -i 's#'"$hookstring"'#HOOKS="base udev autodetect modconf block encrypt filesystems keyboard fsck"#' /mnt/etc/mkinitcpio.conf
        sed -i 's#GRUB_CMDLINE_LINUX=""#GRUB_CMDLINE_LINUX="cryptdevice=UUID='$rootuuid':cryptroot"#' /mnt/etc/default/grub
        sed -i 's/GRUB_CMDLINE_LINUX_DEFAULT="/GRUB_CMDLINE_LINUX_DEFAULT="panic=5 /' /etc/default/grub # Fix for CVE-2016-4484
    fi
fi

arch-chroot /mnt mkinitcpio -p linux
if [ $? -ne 0 ]; then header ERROR: Cannot continue; exit 1; fi
if [ "$nodiskbuild" -eq 0 ]; then
    arch-chroot /mnt grub-install --target=i386-pc "$disk"
    if [ $? -ne 0 ]; then header ERROR: Cannot continue; exit 1; fi
fi
arch-chroot /mnt grub-mkconfig -o /boot/grub/grub.cfg
if [ $? -ne 0 ]; then header ERROR: Cannot continue; exit 1; fi

header Set networking
arch-chroot /mnt systemctl enable openntpd
arch-chroot /mnt systemctl enable netctl
export interface=$(ip link list | grep "state" | cut -f 2 -d ":" | cut -f 2 -d " " | grep -v lo)
cp /mnt/etc/netctl/examples/ethernet-dhcp /mnt/etc/netctl/$interface
sed -i 's/eth0/'$interface'/' /mnt/etc/netctl/$interface
echo 'DNSSearch="aninix.net"' >> /mnt/etc/netctl/$interface
arch-chroot /mnt systemctl enable netctl
arch-chroot /mnt netctl enable $interface

# Vim cleanup for SSH
arch-chroot /mnt mkdir -p /usr/share/vim/vimfiles/plugin
arch-chroot /mnt printf 'set mouse-=a\n' > /usr/share/vim/vimfiles/plugin/shadowarch.vim

ln -sf /etc/skel/.bashrc /mnt/root/.bashrc

# Clone ConfigPackags from AniNIX::Foundation
arch-chroot /mnt git -C /usr/local/src/ clone https://aninix.net/foundation/ConfigPackages
arch-chroot /mnt git -C /usr/local/src/ clone https://aninix.net/foundation/MiscScripts

arch-chroot /mnt make -C /usr/local/src/MiscScripts/Shared install
arch-chroot /mnt make -C /usr/local/src/MiscScripts/Admin install
arch-chroot /mnt make -C /usr/local/src/MiscScripts/ShadowArch install
arch-chroot /mnt git -C /usr/local/src/ clone https://aur.archlinux.org/cower.git
arch-chroot /mnt useradd -m depriv

# Hook for Heartbeat
arch-chroot /mnt /bin/bash -c "mkdir /usr/local/etc/Heartbeat/; echo \"ShadowArch ; /bin/bash -c \\\"systemctl status | grep -c 'State: running'\\\" ; 3\" >> /usr/local/etc/Heartbeat/services.list"

# Handle AUR Packages

if [ "$kali" -eq 1 ]; then
    arch-chroot /mnt git -C /usr/local/src/ clone https://aur.archlinux.org/autopsy.git
fi

# Set password
header Set new root passphrase and depriviledged user '(depriv)' password.
arch-chroot /mnt passwd
arch-chroot /mnt passwd depriv
arch-chroot /mnt chown -R depriv:depriv /usr/local/src/

# Set SSH host keys
arch-chroot /mnt ssh-keygen -A

cp /root/shadowarch /mnt/root/shadowarch.installer."$(date +%F-%R)"

if [ "$gui" -eq 1 ]; then
    echo "Remember to install your graphics drivers!
    For NVidia, look at xf86-video-nouveau
    For AMD, look at xf86-video-amdgpu
    For Hyper-V, look at xf86-video-fbdev
    For Virtual Box, look at virtualbox-guest-utils
    For VMware, look at open-vm-tools"
fi

# Set hostname
header Set hostname
printf "What is your hostname? AniNIX::"
read hostname
echo "$hostname" > /mnt/etc/hostname

header Installed ShadowArch\! 
if [ "$nodiskbuild" -eq 1 ]; then
    header Remember to run grub-install and set up your bootloader.
    echo 'https://wiki.archlinux.org/index.php/Installation_guide#Boot_loader'
else 
    header Press enter to reboot.
    read

    # Reboot
    shutdown -r now
fi