#!/usr/bin/env bash
###############################################################################
#  NIST 800-53 / CIS Partition Layout – Ubuntu 24.04 (Proxmox VM Template)
#
#  PURPOSE:
#    Migrates a default single-partition Ubuntu install into a proper
#    LVM-based layout with separate volumes for /tmp, /home, /var,
#    /var/tmp, /var/log, and /var/log/audit, each with restrictive
#    mount options.
#
#  PREREQUISITES:
#    - A second virtual disk attached to the VM (e.g., /dev/sdb)
#      Add it in Proxmox: Hardware → Add → Hard Disk
#    - Run BEFORE the hardening script
#    - Run from a root shell (sudo -i)
#
#  STRATEGY:
#    Uses LVM on the second disk so volumes can be resized later.
#    After this script, the original root disk holds only / and /boot.
#
#  Usage:
#    chmod +x partition-layout.sh
#    sudo ./partition-layout.sh /dev/sdb      # <-- your second disk
###############################################################################
set -euo pipefail

LOG="/var/log/partition-layout-$(date +%Y%m%d-%H%M%S).log"
exec > >(tee -a "$LOG") 2>&1

banner() { printf '\n\e[1;36m>>> %s\e[0m\n' "$1"; }

if [[ $EUID -ne 0 ]]; then echo "Run as root." >&2; exit 1; fi

DISK="${1:-}"
if [[ -z "$DISK" ]]; then
  echo "Usage: $0 /dev/sdX"
  echo "  Supply the SECOND disk device (the one you added in Proxmox)."
  echo ""
  echo "Available disks:"
  lsblk -d -o NAME,SIZE,TYPE,MOUNTPOINTS | grep disk
  exit 1
fi

if [[ ! -b "$DISK" ]]; then
  echo "ERROR: $DISK is not a valid block device." >&2
  exit 1
fi

# Safety check — abort if the disk has existing partitions with data
if lsblk -n "$DISK" | grep -q 'part'; then
  echo "WARNING: $DISK already has partitions:"
  lsblk "$DISK"
  read -rp "This will DESTROY all data on $DISK. Continue? [y/N] " confirm
  [[ "$confirm" =~ ^[Yy]$ ]] || exit 0
fi

###############################################################################
# CONFIGURATION – Adjust sizes to your needs
# These are minimums; LVM lets you grow them later with lvextend + resize2fs
###############################################################################
VG_NAME="vg_nist"

# Size allocations (adjust to your total disk size)
# For a 50 GB second disk, this uses ~38 GB and leaves ~12 GB free in the VG
declare -A LV_SIZES=(
  [lv_home]="8G"
  [lv_tmp]="2G"
  [lv_var]="10G"
  [lv_var_tmp]="2G"
  [lv_var_log]="8G"
  [lv_var_log_audit]="4G"
)

# Mount options per NIST 800-53 / CIS
declare -A MOUNT_OPTS=(
  [/home]="defaults,nosuid,nodev"
  [/tmp]="defaults,nosuid,nodev,noexec"
  [/var]="defaults,nosuid"
  [/var/tmp]="defaults,nosuid,nodev,noexec"
  [/var/log]="defaults,nosuid,nodev,noexec"
  [/var/log/audit]="defaults,nosuid,nodev,noexec"
)

# Map LV names to mount points
declare -A LV_MOUNTS=(
  [lv_home]="/home"
  [lv_tmp]="/tmp"
  [lv_var]="/var"
  [lv_var_tmp]="/var/tmp"
  [lv_var_log]="/var/log"
  [lv_var_log_audit]="/var/log/audit"
)

###############################################################################
# STEP 1: Install LVM tools if missing
###############################################################################
banner "Installing LVM2"
apt-get update -y && apt-get install -y lvm2

###############################################################################
# STEP 2: Partition the disk and create the physical volume + volume group
###############################################################################
banner "Preparing $DISK for LVM"

wipefs -a "$DISK"
parted -s "$DISK" mklabel gpt
parted -s "$DISK" mkpart primary 1MiB 100%
parted -s "$DISK" set 1 lvm on

# Refresh partition table
partprobe "$DISK"
sleep 2

# Identify the partition (sdb1 or vdb1 etc.)
PART="${DISK}1"
# Handle NVMe naming (nvme0n1p1)
if [[ "$DISK" == *nvme* ]] || [[ "$DISK" == *loop* ]]; then
  PART="${DISK}p1"
fi

pvcreate -f "$PART"
vgcreate "$VG_NAME" "$PART"

echo "Volume group created:"
vgdisplay "$VG_NAME" | grep -E 'VG Name|VG Size|Free'

###############################################################################
# STEP 3: Create logical volumes
###############################################################################
banner "Creating logical volumes"

# Order matters: /var must be populated before /var/log and /var/log/audit
ORDERED_LVS=(lv_var lv_var_log lv_var_log_audit lv_home lv_tmp lv_var_tmp)

for lv in "${ORDERED_LVS[@]}"; do
  lvcreate -L "${LV_SIZES[$lv]}" -n "$lv" "$VG_NAME" -y
  mkfs.ext4 -L "$lv" "/dev/${VG_NAME}/${lv}"
  echo "  Created $lv (${LV_SIZES[$lv]})"
done

echo ""
echo "Remaining free space in VG (available for future growth):"
vgdisplay "$VG_NAME" | grep 'Free'

###############################################################################
# STEP 4: Migrate data from existing directories to new volumes
#
#  We process them in dependency order: /var first, then /var/log,
#  then /var/log/audit, so nested mounts copy correctly.
###############################################################################
banner "Migrating data to new volumes"

# Migration order: deepest-nesting last when copying, first when mounting
MIGRATE_ORDER=(lv_var lv_var_log lv_var_log_audit lv_home lv_tmp lv_var_tmp)

for lv in "${MIGRATE_ORDER[@]}"; do
  mount_point="${LV_MOUNTS[$lv]}"
  dev="/dev/${VG_NAME}/${lv}"
  temp_mount="/mnt/migrate_${lv}"

  echo "Migrating ${mount_point} ..."

  mkdir -p "$temp_mount"
  mount "$dev" "$temp_mount"

  # Copy existing data preserving permissions
  if [[ -d "$mount_point" ]]; then
    rsync -aAXq "${mount_point}/" "$temp_mount/" 2>/dev/null || true
  fi

  umount "$temp_mount"
  rmdir "$temp_mount"

  echo "  ✓ ${mount_point} data copied to ${dev}"
done

###############################################################################
# STEP 5: Build fstab entries
###############################################################################
banner "Updating /etc/fstab"

cp /etc/fstab /etc/fstab.bak.$(date +%s)

# Remove any existing entries for these mount points (avoid duplicates)
for lv in "${ORDERED_LVS[@]}"; do
  mp="${LV_MOUNTS[$lv]}"
  # Escape slashes for sed
  mp_escaped=$(echo "$mp" | sed 's|/|\\/|g')
  sed -i "\| ${mp_escaped} |d" /etc/fstab
done

# Also remove existing tmpfs /tmp entry if the hardening script added one
sed -i '/tmpfs.*\/tmp/d' /etc/fstab

# Append new entries
{
  echo ""
  echo "# ===== NIST 800-53 Separate Partitions (LVM on $DISK) ====="
  for lv in "${ORDERED_LVS[@]}"; do
    mp="${LV_MOUNTS[$lv]}"
    opts="${MOUNT_OPTS[$mp]}"
    echo "/dev/${VG_NAME}/${lv}  ${mp}  ext4  ${opts}  0 2"
  done
  echo ""
  echo "# /dev/shm hardening"
  echo "tmpfs  /dev/shm  tmpfs  defaults,nosuid,nodev,noexec  0 0"
} >> /etc/fstab

echo "New fstab:"
cat /etc/fstab

###############################################################################
# STEP 6: Mount everything and verify
###############################################################################
banner "Mounting all new volumes"

# Must mount in order: /var before /var/log before /var/log/audit
MOUNT_ORDER=(/var /var/log /var/log/audit /home /tmp /var/tmp)

for mp in "${MOUNT_ORDER[@]}"; do
  mount "$mp"
  echo "  ✓ Mounted $mp"
done

# Set sticky bit on /tmp and /var/tmp
chmod 1777 /tmp /var/tmp

# Ensure /var/log/audit exists with correct ownership
mkdir -p /var/log/audit
chmod 700 /var/log/audit

# Remount /dev/shm
mount -o remount /dev/shm 2>/dev/null || true

###############################################################################
# STEP 7: Restart critical services that use /var/log and /tmp
###############################################################################
banner "Restarting services"

systemctl restart rsyslog 2>/dev/null || true
systemctl restart systemd-journald 2>/dev/null || true
systemctl restart auditd 2>/dev/null || true

###############################################################################
# VERIFICATION
###############################################################################
banner "Verification"

echo ""
echo "━━━ Mount Points ━━━"
df -h /home /tmp /var /var/tmp /var/log /var/log/audit 2>/dev/null | \
  awk 'NR==1 || /vg_nist/'
echo ""
echo "━━━ Mount Options ━━━"
for mp in "${MOUNT_ORDER[@]}"; do
  opts=$(findmnt -n -o OPTIONS "$mp" 2>/dev/null || echo "NOT MOUNTED")
  printf "  %-20s %s\n" "$mp" "$opts"
done

echo ""
echo "━━━ LVM Free Space (for future growth) ━━━"
vgdisplay "$VG_NAME" | grep -E 'Free|VG Size'

cat <<'DONE'

╔══════════════════════════════════════════════════════════════════╗
║             Partition Layout Complete                            ║
╠══════════════════════════════════════════════════════════════════╣
║                                                                  ║
║  Separate mount points created:                                  ║
║    /home           – nosuid, nodev                               ║
║    /tmp            – nosuid, nodev, noexec                       ║
║    /var            – nosuid                                      ║
║    /var/tmp        – nosuid, nodev, noexec                       ║
║    /var/log        – nosuid, nodev, noexec                       ║
║    /var/log/audit  – nosuid, nodev, noexec                       ║
║    /dev/shm        – nosuid, nodev, noexec                       ║
║                                                                  ║
║  To grow a volume later:                                         ║
║    lvextend -L +5G /dev/vg_nist/lv_var_log                       ║
║    resize2fs /dev/vg_nist/lv_var_log                             ║
║                                                                  ║
║  NEXT: Reboot to confirm everything mounts cleanly,              ║
║        then run the hardening script.                            ║
╚══════════════════════════════════════════════════════════════════╝
DONE

echo "Full log: $LOG"