#!/bin/bash

function header () {
    tput setaf 1
    tput bold
    echo $@
    tput sgr0
    return
}
function help() {
    echo Usage: ${0} '[OPTIONS]'
    echo '\-A      -- Audio optimizations from the Arch Wiki'
    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 '\-l FILE -- Log to a file'
    echo '\-p      -- Productivity package additions'
    echo '\-P      -- Power saving for laptops'
    echo '\-s      -- Create a layout for an AniNIX::Spartacus'
    echo '\-m      -- Skip disk operations and assume storage is mounted on /mnt'
    echo '\-v      -- Verbose output.'
    echo '\-z      -- Try to add all the packages on AniNIX::Core'
    exit 1;
}

# Partition controls
efipart=2;
bootpart=3;
rootpart=4;
partpoint=1;
partedcmd='mklabel gpt\nmkpart primary ext2 0 1MiB\nset 1 bios_grub on\n';
function addNextPartition() {
    partsize="$1"
    parttype="$2"
    partfs="$3"
    nextpartpoint=$(( $partpoint + $partsize ))
    partedcmd="${partedcmd}mkpart $parttype $partfs ${partpoint}MiB ${nextpartpoint}MiB"'\n'
    partpoint=$nextpartpoint
}

disk="/dev/sda"
bootsize=500; # Size in MB for /boot
# TODO Add LVM as an argument
while getopts "d:egkl:pmsvz" OPTION
do
    case $OPTION in
        A) audio=1 ;;
        d) disk=${OPTARG} ;;
        e) encrypt=1 ;;
        g) gui=1 ;;
        k) kali=1 ;;
        l) exec script -e -f -c "/bin/bash $0 $(echo $@ | sed "s#-l ${OPTARG}##")" "${OPTARG}" ;;
        p) productivity=1; gui=1 ;;
        P) powersave=1 ;;
        m) nodiskbuild=1 ;;
        s) spartacus=1 ;;
        v) set -x ;;
        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 [ -z "$nodiskbuild" ]; then
    header Allocating space
    dd if=/dev/zero of="$disk" bs=1M count=1000

    if [ ! -z "$spartacus" ]; then
        # Insert an ExFAT data partition ahead of the rest.
        export datapart=$efipart;
        export efipart=$((efipart+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))
        addNextPartition $datasize primary ext4
    fi
    # 550MiB for EFI with boot toggle
    addNextPartition 550 primary fat32
    partedcmd="${partedcmd}toggle $efipart boot"'\n'

    # /boot
    addNextPartition $bootsize primary fat32

    # / (root)
    partedcmd="${partedcmd}mkpart primary ext4 ${partpoint}MiB 100%%FREE"'\nquit\n\n'
    printf "$partedcmd" | parted "$disk"
    if [ ! -z "$spartacus" ]; then
        #create data partition
        pacman -S exfat-utils --noconfirm
        mkfs.exfat "$disk""$datapart"
        exfatlabel "$disk""$datapart" "AS-XPLATFRM"
    fi

    header Making fat esp partition on "$disk""$efipart"
    mkfs.fat -F32 "$disk""$efipart"

    header Making boot partition on "$disk""$bootpart"
    mkfs.vfat "$disk""$bootpart"

    header Making root and mountpoints
    if [ ! -z "$encrypt" ]; 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
    mkdir /mnt/boot/efi
    mount "$disk""$efipart" /mnt/boot/efi
    if [ "$?" -ne 0 ]; then header ERROR: Cannot continue; exit 1; fi
fi


# Install ArchLinux with basic clients for the AniNIX Services.
# * git for Foundation
# * elinks for WebServer and Wiki
# * openssh for SSH/SFTP
# * weechat 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 elinks weechat vim wget tor torsocks grub os-prober rsync openntpd tmux efibootmgr"
if [ ! -z "$gui" ]; then 
    export pkglist="$pkglist"" xorg-server xfce4 chromium conky tigervnc xscreensaver"
fi
if [ ! -z "$spartacus" ]; then
    export pkglist="$pkglist"" exfat-utils"
fi
if [ ! -z "$productivity" ]; then
    export pkglist="$pkglist"" libreoffice-still gimp feh vlc evince openshot"
fi
if [ ! -z "$kali" ]; then
    export pkglist="$pkglist"" extundelete testdisk nmap tcpdump hexedit dcfldd"
    if [ ! -z "$gui" ]; then
        export pkglist="$pkglist"" wireshark-gtk"
    else 
        export pkglist="$pkglist"" wireshark-cli"
    fi
fi
if [ ! -z "$kitchensink" ]; 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 [ -z "$nodiskbuild" ]; then
    export rootuuid="$(blkid "$disk""$rootpart" | cut -f 2 -d '"')"
    if [ ! -z "$encrypt" ]; 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 [ -z "$nodiskbuild" ]; then
    arch-chroot /mnt grub-install --target=x86_64-efi --removable --bootloader-id=grub --efi-directory /boot "$disk"
    if [ $? -ne 0 ]; then header ERROR: Cannot continue; exit 1; fi
    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 groupadd tty-allow
arch-chroot /mnt useradd -u 1001 -G tty-allow -m depriv
arch-chroot /mnt usermod -G "$(getent group | grep root | cut -f 1 -d ':' | tr '\n' ',')""tty-allow" root
arch-chroot /mnt /bin/bash -c 'line="$(grep -E root"[[:space:]]"ALL /etc/sudoers)"; sed -i "s/$line/$line\ndepriv ALL=(ALL) ALL/" /etc/sudoers'

# Handle AUR Packages

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

# Optimizations from https://wiki.archlinux.org/index.php/Power_management
if [ ! -z "$powersave" ]; then
    if [ `lspci | grep -i intel | grep -ic audio` -eq 1 ]; then
        echo 'options snd_hda_intel power_save=1' > /mnt/etc/modprobe.d/audio_powersave.conf
    else 
        echo 'options snd_ac97_codec power_save=1' > /mnt/etc/modprobe.d/audio_powersave.conf
    fi
    arch-chroot /mnt pacman -S rfkill cpupower --noconfirm
    arch-chroot /mnt systemctl enable rfkill-block@.service
    echo 'kernel.nmi_watchdog = 0' > /mnt/etc/sysctl.d/disable_watchdog.conf
    echo 'vm.dirty_writeback_centisecs = 6000' > /mnt/etc/sysctl.d/dirty_writes.conf
    echo 'vm.laptop_mode = 5' > /mnt/etc/sysctl.d/laptop.conf
    echo 'ACTION=="add", SUBSYSTEM=="net", KERNEL=="wlan*", RUN+="/usr/bin/iw dev %k set power_save on"' > /mnt/etc/udev/rules.d/70-wifi-powersave.rules
    echo 'blacklist uvcvideo' > /mnt/etc/modprobe.d/no-camera.conf
fi

# Thanks to https://wiki.archlinux.org/index.php/Professional_audio
if [ ! -z "$audio" ]; then
    sed -i 's#GRUB_CMDLINE_LINUX_DEFAULT="#GRUB_CMDLINE_LINUX_DEFAULT="threadirqs #' /mnt/etc/default/grub
    arch-chroot /mnt grub-mkconfig -o /boot/grub/grub.cfg
    printf 'vm.swappiness = 10\nfs.inotify.max_user_watches = 524288\n' > /mnt/etc/sysctl.d/99-audio-tuning.conf
    setpci -v -d *:* latency_timer=b0
    for SOUND_CARD_PCI_ID in `lspci | grep -i audio | cut -f 1 -d ' '`; do
        setpci -v -s $SOUND_CARD_PCI_ID latency_timer=ff;
    done
    printf 'echo 2048 > /sys/class/rtc/rtc0/max_user_freq\necho 2048 > /proc/sys/dev/hpet/max-user-freq\n' >> /mnt/etc/rc.local
fi

# Emulate old rc.local behavior
touch /mnt/etc/rc.local
echo '[Unit]
Description=Execute commands at boot

[Service]
ExecStart=/etc/rc.local

[Install]
WantedBy=multi-user.target' > /mnt/usr/lib/systemd/system/rclocal.service
arch-chroot /mnt systemctl enable rclocal
arch-chroot /mnt chmod 0700 /etc/rc.local
arch-chroot /mnt chown root: /etc/rc.local

# 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 [ ! -z "$gui" ]; 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 [ ! -z "$nodiskbuild" ]; 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