#!/bin/sh
# Part of raspi-config http://github.com/asb/raspi-config
#
# See LICENSE file for copyright and license details

ASK_TO_REBOOT=0

do_info() {
  whiptail --msgbox "\
This tool provides a straight-forward way of doing initial 
configuration of the Raspberry Pi. Although it can be run 
at any time, some of the options may have difficulties if 
you have heavily customised your installation.\
" 20 70 1
}

do_expand_rootfs() {
  # Get the starting offset of the root partition
  PART_START=$(parted /dev/mmcblk0 -ms unit s p | grep "^2" | cut -f 2 -d:)
  [ "$PART_START" ] || return 1
  # Return value will likely be error for fdisk as it fails to reload the 
  # partition table because the root fs is mounted
  fdisk /dev/mmcblk0 <<EOF
p
d
2
n
p
2
$PART_START

p
w
EOF
  ASK_TO_REBOOT=1

  # now set up an init.d script
cat <<\EOF > /etc/init.d/resize2fs_once &&
#!/bin/sh
### BEGIN INIT INFO
# Provides:          resize2fs_once
# Required-Start:
# Required-Stop:
# Default-Start: 2 3 4 5 S
# Default-Stop:
# Short-Description: Resize the root filesystem to fill partition
# Description:
### END INIT INFO

. /lib/lsb/init-functions

case "$1" in
  start)
    log_daemon_msg "Starting resize2fs_once" &&
    resize2fs /dev/mmcblk0p2 &&
    rm /etc/init.d/resize2fs_once &&
    update-rc.d resize2fs_once remove &&
    log_end_msg $?
    ;;
  *)
    echo "Usage: $0 start" >&2
    exit 3
    ;;
esac
EOF
  chmod +x /etc/init.d/resize2fs_once &&
  update-rc.d resize2fs_once defaults &&
  whiptail --msgbox "Root partition has been resized.\n\
The filesystem will be enlarged upon the next reboot" 20 60 2
}

set_config_var() {
  lua - "$1" "$2" "$3" <<EOF > "$3.bak"
local key=assert(arg[1])
local value=assert(arg[2])
local fn=assert(arg[3])
local file=assert(io.open(fn))
local made_change=false
for line in file:lines() do
  if line:match("^#?%s*"..key.."=.*$") then
    line=key.."="..value
    made_change=true
  end
  print(line)
end

if not made_change then
  print(key.."="..value)
end
EOF
mv "$3.bak" "$3"
}

get_config_var() {
  lua - "$1" "$2" <<EOF
local key=assert(arg[1])
local fn=assert(arg[2])
local file=assert(io.open(fn))
for line in file:lines() do
  local val = line:match("^#?%s*"..key.."=(.*)$")
  if (val ~= nil) then
    print(val)
    break
  end
end
EOF
}

# $1 is 0 to disable overscan, 1 to disable it
set_overscan() {
  # Stop if /boot is not a mountpoint
  if ! mountpoint -q /boot; then
    return 1
  fi

  [ -e /boot/config.txt ] || touch /boot/config.txt

  if [ "$1" -eq 0 ]; then # disable overscan
    sed /boot/config.txt -i -e "s/^overscan_/#overscan_/"
    set_config_var disable_overscan 1 /boot/config.txt
  else # enable overscan
    set_config_var disable_overscan 0 /boot/config.txt
  fi
}

do_overscan() {
  whiptail --yesno "What would you like to do with overscan" 20 60 2 \
    --yes-button Disable --no-button Enable 
  RET=$?
  if [ $RET -eq 0 ] || [ $RET -eq 1 ]; then
    ASK_TO_REBOOT=1
    set_overscan $RET;
  else
    return 1
  fi
}

do_change_pass() {
  whiptail --msgbox "You will now be asked to enter a new password for the pi user" 20 60 1
  passwd pi &&
  whiptail --msgbox "Password changed successfully" 20 60 1
}

do_configure_keyboard() {
  dpkg-reconfigure keyboard-configuration &&
  printf "Reloading keymap. This may take a short while\n" &&
  invoke-rc.d keyboard-setup start
}

do_change_locale() {
  dpkg-reconfigure locales
}

do_change_timezone() {
  dpkg-reconfigure tzdata
}

do_memory_split() {
  if [ -e /boot/start_cd.elf ]; then
    # New-style memory split setting
    if ! mountpoint -q /boot; then
      return 1
    fi
    ## get current memory split from /boot/config.txt
    CUR_GPU_MEM=$(get_config_var gpu_mem /boot/config.txt)
    [ -z "$CUR_GPU_MEM" ] && CUR_GPU_MEM=64
    ## ask users what gpu_mem they want
    NEW_GPU_MEM=$(whiptail --inputbox "How much memory should the GPU have?  e.g. 16/32/64/128/256" \
      20 70 "$CUR_GPU_MEM" 3>&1 1>&2 2>&3)
    if [ $? -eq 0 ]; then
      set_config_var gpu_mem "$NEW_GPU_MEM" /boot/config.txt
      ASK_TO_REBOOT=1
    fi
  else # Old firmware so do start.elf renaming
    get_current_memory_split
    MEMSPLIT=$(whiptail --menu "Set memory split.\n$MEMSPLIT_DESCRIPTION" 20 60 10 \
      "240" "240MiB for ARM, 16MiB for VideoCore" \
      "224" "224MiB for ARM, 32MiB for VideoCore" \
      "192" "192MiB for ARM, 64MiB for VideoCore" \
      "128" "128MiB for ARM, 128MiB for VideoCore" \
      3>&1 1>&2 2>&3)
    if [ $? -eq 0 ]; then
      set_memory_split ${MEMSPLIT}
      ASK_TO_REBOOT=1
    fi
  fi
}

get_current_memory_split() {
  # Stop if /boot is not a mountpoint
  if ! mountpoint -q /boot; then
    return 1
  fi
  AVAILABLE_SPLITS="128 192 224 240"
  MEMSPLIT_DESCRIPTION=""
  for SPLIT in $AVAILABLE_SPLITS;do
    if [ -e /boot/arm${SPLIT}_start.elf ] && cmp /boot/arm${SPLIT}_start.elf /boot/start.elf >/dev/null 2>&1;then
      CURRENT_MEMSPLIT=$SPLIT
      MEMSPLIT_DESCRIPTION="Current: ${CURRENT_MEMSPLIT}MiB for ARM, $((256 - $CURRENT_MEMSPLIT))MiB for VideoCore"
      break
    fi
  done
}

set_memory_split() {
  cp -a /boot/arm${1}_start.elf /boot/start.elf
  sync
}

do_overclock() {
  whiptail --msgbox "\
Be aware that overclocking may reduce the lifetime of your
Raspberry Pi. If overclocking at a certain level causes
system instability, try a more modest overclock. See
http://elinux.org/RPi_Overclocking for more information.\
" 20 70 1
  OVERCLOCK=$(whiptail --menu "Chose overclock preset" 20 60 10 \
    "None" "700MHz ARM, 250MHz core, 400MHz SDRAM, 0 overvolt" \
    "Modest" "800MHz ARM, 300MHz core, 400MHz SDRAM, 0 overvolt" \
    "Medium" "900MHz ARM, 333MHz core, 450MHz SDRAM, 2 overvolt" \
    "High" "950MHz ARM, 450MHz core, 450MHz SDRAM, 6 overvolt" \
    "Turbo" "1000MHz ARM, 500MHz core, 500MHz SDRAM, 6 overvolt" \
    3>&1 1>&2 2>&3)
  if [ $? -eq 0 ]; then
    case "$OVERCLOCK" in
      None)
        set_overclock None 700 250 400 0
        ;;
      Modest)
        set_overclock Modest 800 300 400 0
        ;;
      Medium)
        set_overclock Medium 900 333 450 2
        ;;
      High)
        set_overclock High 950 450 450 6
        ;;
      Turbo)
        set_overclock Turbo 1000 500 500 6
        ;;
      *)
        whiptail --msgbox "Programmer error, unrecognised overclock preset" 20 60 2
        return 1
        ;;
    esac
    ASK_TO_REBOOT=1
  fi
}

set_overclock() {
  set_config_var arm_freq $2 /boot/config.txt &&
  set_config_var core_freq $3 /boot/config.txt &&
  set_config_var sdram_freq $4 /boot/config.txt &&
  set_config_var over_voltage $5 /boot/config.txt &&
  # now set up an init.d script
cat <<\EOF > /etc/init.d/switch_cpu_governor &&
#!/bin/sh
### BEGIN INIT INFO
# Provides:          switch_cpu_governor
# Required-Start: udev mountkernfs $remote_fs
# Required-Stop:
# Default-Start: S
# Default-Stop:
# Short-Description: Switch to ondemand cpu governor (unless shift key is pressed)
# Description:
### END INIT INFO

. /lib/lsb/init-functions

case "$1" in
  start)
    log_daemon_msg "Checking if shift key is held down"
    timeout 1 thd --dump /dev/input/event* | grep -q "LEFTSHIFT\|RIGHTSHIFT"
    if [ $? -eq 0 ]; then
      printf " Yes. Not switching scaling governor"
      log_end_msg 0
    else
      SYS_CPUFREQ_GOVERNOR=/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor
      [ -e $SYS_CPUFREQ_GOVERNOR ] && echo "ondemand" > $SYS_CPUFREQ_GOVERNOR
      printf " No. Switching to ondemand scaling governor"
      log_end_msg 0
    fi
    ;;
  *)
    echo "Usage: $0 start" >&2
    exit 3
    ;;
esac
EOF
  chmod +x /etc/init.d/switch_cpu_governor &&
  update-rc.d switch_cpu_governor defaults &&
  whiptail --msgbox "Set overclock to preset '$1'" 20 60 2
}

do_ssh() {
  if [ -e /var/log/regen_ssh_keys.log ] && ! grep -q "^finished" /var/log/regen_ssh_keys.log; then
    whiptail --msgbox "Initial ssh key generation still running. Please wait and try again." 20 60 2
    return 1
  fi
  whiptail --yesno "Would you like the SSH server enabled or disabled?" 20 60 2 \
    --yes-button Enable --no-button Disable 
  RET=$?
  if [ $RET -eq 0 ]; then
    update-rc.d ssh enable &&
    invoke-rc.d ssh start &&
    whiptail --msgbox "SSH server enabled" 20 60 1
  elif [ $RET -eq 1 ]; then
    update-rc.d ssh disable &&
    whiptail --msgbox "SSH server disabled" 20 60 1
  else
    return $RET
  fi
}

do_boot_behaviour() {
  whiptail --yesno "Should we boot straight to desktop?" 20 60 2
  RET=$?
  if [ $RET -eq 0 ]; then # yes
    update-rc.d lightdm enable 2
    sed /etc/lightdm/lightdm.conf -i -e "s/^#autologin-user=.*/autologin-user=pi/"
    ASK_TO_REBOOT=1
  elif [ $RET -eq 1 ]; then # no
    update-rc.d lightdm disable 2
    ASK_TO_REBOOT=1
  else # user hit escape
    return 1
  fi
}

do_update() {
  apt-get update &&
  apt-get install raspi-config &&
  printf "Sleeping 5 seconds before reloading raspi-config\n" &&
  sleep 5 &&
  exec raspi-config
}

do_finish() {
  if [ -e /etc/profile.d/raspi-config.sh ]; then
    rm -f /etc/profile.d/raspi-config.sh
    sed -i /etc/inittab \
      -e "s/^#\(.*\)#\s*RPICFG_TO_ENABLE\s*/\1/" \
      -e "/#\s*RPICFG_TO_DISABLE/d"
    telinit q
  fi
  if [ $ASK_TO_REBOOT -eq 1 ]; then
    whiptail --yesno "Would you like to reboot now?" 20 60 2
    if [ $? -eq 0 ]; then # yes
      sync
      reboot
    fi
  fi
  exit 0
}

#
# Command line options for non-interactive use
#
for i in $*
do
  case $i in
  --memory-split)
    OPT_MEMORY_SPLIT=GET
    ;;
  --memory-split=*)
    OPT_MEMORY_SPLIT=`echo $i | sed 's/[-a-zA-Z0-9]*=//'`
    ;;
  *)
    # unknown option
    ;;
  esac
done

if [ "GET" = "${OPT_MEMORY_SPLIT:-}" ]; then
  set -u # Fail on unset variables
  get_current_memory_split
  echo $CURRENT_MEMSPLIT
  exit 0
fi

# Everything else needs to be run as root
if [ $(id -u) -ne 0 ]; then
  printf "Script must be run as root. Try 'sudo raspi-config'\n"
  exit 1
fi

if [ -n "${OPT_MEMORY_SPLIT:-}" ]; then
  set -e # Fail when a command errors
  set_memory_split "${OPT_MEMORY_SPLIT}"
  exit 0
fi

#
# Interactive use loop
#
while true; do
  FUN=$(whiptail --menu "Raspi-config" 20 80 12 --cancel-button Finish --ok-button Select \
    "info" "Information about this tool" \
    "expand_rootfs" "Expand root partition to fill SD card" \
    "overscan" "Change overscan" \
    "configure_keyboard" "Set keyboard layout" \
    "change_pass" "Change password for 'pi' user" \
    "change_locale" "Set locale" \
    "change_timezone" "Set timezone" \
    "memory_split" "Change memory split" \
    "overclock" "Configure overclocking" \
    "ssh" "Enable or disable ssh server" \
    "boot_behaviour" "Start desktop on boot?" \
    "update" "Try to upgrade raspi-config" \
    3>&1 1>&2 2>&3)
  RET=$?
  if [ $RET -eq 1 ]; then
    do_finish
  elif [ $RET -eq 0 ]; then
    "do_$FUN" || whiptail --msgbox "There was an error running do_$FUN" 20 60 1
  else
    exit 1
  fi
done