#!/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

#Argument check!

while getopts ud:r: option
do
    case "${option}" in
        u) UNATTEND=true;;
        d) DESKTOP=${OPTARG};;
        r) REPO=${OPTARG};;
    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

DISTRO=$(lsb_release -i -s)

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

case $DISTRO in
Kali)
    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/02-unattended-boot.binary ./build/live-build-config/kali-config/common/includes.installer/
    cp config/preseed.cfg ./build/live-build-config/kali-config/common/includes.installer/
    cp config/software.list.chroot ./build/live-build-config/kali-config/variant-$DESKTOP/package-lists/

    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
    cp images/*.iso ../../images/kali-rolling-$DESKTOP-$(date -I).iso
;;
Parrot)
    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/

    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
    cp ../*.iso ../../images/parrotsec-$DESKTOP-$(date -I).iso
;;
esac

if [ $UNATTEND != "true" ];
then
whiptail --backtitle "Pentest Build Script" --title "Finished" --msgbox "Fingers crossed there were no errors... Your image should be done!" 7 70
fi

exit 0
