#!/bin/bash

function pause(){
    read -p "$*"
}

#Root check!
if [ "$EUID" -ne 0 ] 
then
    whiptail --backtitle "Pentest Build Script" --title "NOT ROOT" --msgbox "YOU ARE NOT ROOT.  PLEASE RUN THIS SCRIPT AS ROOT." 10 55
    exit 1
fi

#Arguments

DISTRO=$(lsb_release -i -s)

while getopts ul:d:r: option
do
    case "${option}" in
        u) UNATTEND="true" && echo "UNATTEND";;
        l) DISTRO=${OPTARG} && echo $DISTRO;;
        d) DESKTOP=${OPTARG} && echo $DESKTOP;;
        r) REPO=${OPTARG} && echo $REPO;;
    esac
done


if [[ $UNATTEND != "true" ]];
then
    whiptail --backtitle "Pentest Build Script" --title "Welcome" --msgbox "This script will help you build a custom pentesting distribution." 7 70

    DESKTOP=$(whiptail --backtitle "Pentest Build Script" --title "Desktop Environment" --menu "Choose a desktop:" 15 50 6 \
    "gnome" "GNOME Desktop Environment" \
    "mate"  "MATE Advanced Traditional Environment" \
    "cinnamon"  "Cinnamon Desktop Environment" \
    "kde" "K Desktop Environment" \
    "xfce" "Xfce Desktop Environment" \
    "lxde" "Lightweight X11 Desktop Environment" 3>&1 1>&2 2>&3)

    REPO=$(whiptail --backtitle "Pentest Build Script" --title "Repository" --menu "Which repository are you going to use?" 10 60 2 \
    "default" "The distribution's default repository" \
    "local" "A repository hosted on your local network" 3>&1 1>&2 2>&3)

    if [ $REPO == "local" ];
    then
        REPO=$(whiptail --backtitle "Pentest Build Script" --title "Local Repository" --inputbox "What is the IP address of your local repository?" \
        10 60 3>&1 1>&2 2>&3)
    fi

    if (whiptail --backtitle "Pentest Build Script" --title "Confirmation" --yesno "We're going to build:
    Distribution: $DISTRO
    Desktop Environment: $DESKTOP
    Repository: $REPO
    Shall we proceed?" 10 50) 
    then
        whiptail --backtitle "Pentest Build Script" --title "Build Time!" --msgbox "This will take a while.  Press OK to proceed." 8 30
    else
        whiptail --backtitle "Pentest Build Script" --title "OK" --msgbox "This script will now terminate." 8 35
        exit 1
    fi
fi

#This is the stuff that's common to both distros
dpkg-name ./debs/*.deb
rm -rf build
mkdir build
apt install git build-essential cdebootstrap live-build curl wget -y

case $DISTRO in
Kali)
    IMAGENAME=kali-rolling-$DESKTOP-$(date -I)-$(date +%H%M)
    cp -rv ./modules/live-build-config ./build/
    mkdir ./build/live-build-config/kali-config/variant-$DESKTOP/packages.chroot
    cp ./debs/*.deb build/live-build-config/kali-config/variant-$DESKTOP/packages.chroot/
    cp -rv filesystem/* ./build/live-build-config/kali-config/common/includes.chroot/
    cp config/software.list.chroot ./build/live-build-config/kali-config/variant-$DESKTOP/package-lists/
    cp config/pentest-distro-builder.chroot ./build/live-build-config/kali-config/common/hooks/live/

    if [ $REPO != "default" ];
    then
        sed -i "s|archive.kali.org|$REPO|g" ./build/live-build-config/auto/config
    fi

    cd build/live-build-config
    ./build.sh --distribution kali-rolling --variant $DESKTOP --verbose | tee ../../logs/$IMAGENAME.log
    cp images/*.iso ../../images/$IMAGENAME.iso
;;
Parrot)
    IMAGENAME=parrotsec-$DESKTOP-$(date -I)-$(date +%H%M)
    cp -rv ./modules/parrot-build ./build/
    mkdir ./build/parrot-build/templates/parrot-variant-$DESKTOP/packages.chroot
    cp ./debs/*.deb build/parrot-build/templates/parrot-variant-$DESKTOP/packages.chroot/
    cp -rv filesystem/* ./build/parrot-build/templates/common/includes.chroot/
    cp config/software.list.chroot ./build/parrot-build/templates/parrot-variant-$DESKTOP/package-lists/
    cp config/pentest-distro-builder.chroot ./build/parrot-build/templates/common/hooks/normal/

    if [ $REPO != "default" ];
    then
        sed -i "s|archive.parrotsec.org|$REPO|g" ./build/parrot-build/auto/config
        sed -i "s|deb.parrotsec.org|$REPO|g" ./build/parrot-build/templates/common/archives/parrot.list
    fi

    cd build/parrot-build
    ./build.sh build variant-$DESKTOP amd64 | tee ../../logs/$IMAGENAME.log
    cp ../*.iso ../../images/$IMAGENAME.iso
;;
esac

if [[ $UNATTEND != "true" ]];
then
    if [ -f  ../../images/$IMAGENAME.iso ];
    then
        whiptail --backtitle "Pentest Build Script" --title "Finished" \
        --msgbox "Fingers crossed there were no errors... Your image should be done!" 7 70
    else
        whiptail --backtitle "Pentest Build Script" --title "OH NO" \
        --msgbox "Something went wrong and no .iso file was created... Check the log file for details! :-(" 7 92
    fi
fi

exit 0