#!/bin/bash
# apt-rollback Script
# By Fabio Dell'Aria - fabio.dellaria@gmail.com - Mar 2020

# Check if the current user is "root" otherwise restart the script with "sudo"...
[ "$(whoami)" != "root" ] && exec sudo -- "$0" "$@"

# Main Variables...
# --------------------------------------------------------------------------------
{
  MAX_CACHE_SIZE="2000" # In MBytes
  LOG_FILE="/var/log/apt/history.log"
  PKG_DIR="/var/lib/apt-rollback/packages"
  TMP_FILE="/tmp/dpkg.apt-rollback.log"
  VERSION="1.0.16" # Use always the 'x.y.z' format
  UBUNTU_NAME=$(cat /etc/os-release | grep "UBUNTU_CODENAME" | cut -d'=' -f2)
  ARCHITECTURE=$(dpkg --print-architecture)
  REPOSITORY="https://launchpad.net/ubuntu/$UBUNTU_NAME/$ARCHITECTURE/"
  VERSION_FILE="/var/log/installer/$(ls -t /var/log/installer/ 2>/dev/null | head -n 1)"
  INSTALLATION_DATE=$(date -r $VERSION_FILE "+%Y-%m-%d  %H:%M:%S" 2>/dev/null)
  RESULT=""
  DEFAULT_IFS="$IFS"
}
# --------------------------------------------------------------------------------

mkdir -p "$PKG_DIR"

# Change the APT Cache configuration to allow more .DEB file storing (and Downgrading)
{
  echo 'APT::Archives::MaxAge "0";'
  echo 'APT::Archives::MinAge "0";'
  echo 'APT::Archives::MaxSize "2000";'
} > /etc/apt/apt.conf.d/99apt-rollback

# Main Functions...
# --------------------------------------------------------------------------------
function MaintainCache ()
{
  SAVED_IFS="$IFS"
  IFS="$DEFAULT_IFS"
  MBYTE=$((1024*1024))
  CACHE_SIZE=$((MAX_CACHE_SIZE*MBYTE))
  DIR_SIZE=$(du -s "$PKG_DIR" 2>/dev/null | grep -o -E "[0-9]+")
  DIR_SIZE=$((DIR_SIZE*1024))
  PACKAGES=$(ls -rt "$PKG_DIR")
  for FILE_NAME in $PACKAGES; do
    if [ "$DIR_SIZE" -gt "$CACHE_SIZE" ]; then
      FILE_NAME="$PKG_DIR/$FILE_NAME"
      FILE_SIZE=$(stat -c "%s" "$FILE_NAME")
      DIR_SIZE=$((DIR_SIZE-FILE_SIZE))
      rm -f "$FILE_NAME"
    else
      break
    fi
  done
  IFS="$SAVED_IFS"
}

function Progress_Bar ()
{
  while ps | grep $! &>/dev/null; do
    echo -n -e "\e[1;32m.\e[0m"
    sleep 0.5
  done
}

function Undo_Last_Command()
{
  # Check '--last' parameter value
  {
    if [ "$1" == "" ]; then
      LAST="1"
    else
      LAST="$1"
    fi
    INT_REG='^[0-9]+$'
    if ! [[ $LAST =~ $INT_REG ]]; then
      echo "You have to enter a valid integer value for the '--last' parameter"
      exit 1
    fi
  }

  # Set the Words to use (singular, plural)
  {
    if [ "$LAST" == "1" ]; then
      COUNT=""
      PLUS=""
      PRONOUN="it"
    else
      COUNT=" $LAST"
      PLUS="s"
      PRONOUN="them"
    fi
  }

  # Get the last Compressed and Plain-Text Log APT commands
  {
    if [ -f "$LOG_FILE.1.gz" ]; then
      COMMAND_GZIP=$(zcat "$LOG_FILE.1.gz")"\n"
    else
      COMMAND_GZIP=""
    fi
    COMMAND_PLAIN=$(cat "$LOG_FILE")
    COMMANDS_LIST=$(echo "$COMMAND_GZIP$COMMAND_PLAIN" | grep -A2 -E "^(Install|Remove|Purge|Downgrade|Upgrade): " | grep -v -E "^(Commandline|Requested-By|Start-Date|--|$)" | tail -"$((LAST*2))" | tac)
  }

  # List last <n> APT commands
  {
    PKG_DATE=""
    echo "The last$COUNT APT command$PLUS, performed the following Package$PLUS operation$PLUS..."
    echo
    N=0
    IFS=$'\n'
    for COMMAND_LINE in $COMMANDS_LIST; do

      if [[ $COMMAND_LINE = End-Date* ]]; then
        PKG_DATE=$(echo "$COMMAND_LINE" | cut -d' ' -f2-) 
        continue;
      fi

      if [[ "$PKG_DATE" < "$INSTALLATION_DATE" ]]; then
        continue;
      fi

      let "N=N+1"
      LINE=$(echo -e "$COMMAND_LINE" | sed "s/[(][^)]*[)]//g" | sed "s/ ,//g" | sed 's/ *$//g')
      COMMAND=$(echo "$LINE" | cut -d":" -f1)
      PACKAGES=$(echo "$LINE" | cut -d" " -f2-)
      echo -e -n "\e[1;34m#$N\e[0m: "
      FMT_DATE=$(echo "$PKG_DATE" | sed 's/  / /')
      echo -e "$FMT_DATE - $COMMAND of \e[1;32m$PACKAGES\e[0m"
    done
    echo
  }

  Ask_Yes_No "Do you wish to Undo $PRONOUN?"
  if [ "$RESULT" == "y" ]; then
    N=1
    PKG_DATE=""
    for COMMAND_LINE in $COMMANDS_LIST; do

      if [[ $COMMAND_LINE = End-Date* ]]; then 
        PKG_DATE=$(echo "$COMMAND_LINE" | cut -d' ' -f2-) 
        continue;
      fi

      if [[ "$PKG_DATE" < "$INSTALLATION_DATE" ]]; then
        continue;
      fi

      # Showing Command to UNDO...
      {
        LINE=$(echo -e "$COMMAND_LINE" | sed "s/[(][^)]*[)]//g" | sed "s/ ,//g" | sed 's/ *$//g')
        COMMAND=$(echo "$LINE" | cut -d":" -f1)
        PACKAGES=$(echo "$LINE" | cut -d" " -f2-)
        echo -e -n "\e[1;34m#$N\e[0m "
        echo -e "UNDOING: $COMMAND of \e[1;32m${PACKAGES:0:40}\e[0m ..."
      }

      case $COMMAND in
        "Install") COMMAND="purge" ;;
        "Purge" | "Remove") COMMAND="install" ;;
      esac

      # UNDO Install/Remove/Purge...
      if [ "$COMMAND" == "install" ] || [ "$COMMAND" == "purge" ]; then
        IFS=" "
        apt-get $COMMAND -y -qq $PACKAGES >/dev/null &
        PID=$!
        Progress_Bar
        wait $PID
        EXIT_CODE=$?
        IFS=$'\n'
        echo
        if [ ! "$EXIT_CODE" == "0" ]; then
          exit 1
        fi
      else
        # UNDO Upgrade/Downgrade...
        IFS=')'
        DPKG_COMMAND=""
        LOG_COMMAND=""
        COMMANDS=$(echo "$COMMAND_LINE" | cut -d" " -f2-)
        
        for SINGLE_PACKAGE in $COMMANDS; do
          # Remove optional initial comma and space
          SINGLE_PACKAGE=$(echo "$SINGLE_PACKAGE" | sed 's/^, //') 

          # Get Single Package NAME, ARCHITECTURE, OLD_VERSION and NEW_VERSION...
          {
            PACKAGE_NAME=$(echo "$SINGLE_PACKAGE" | cut -d':' -f1)
            PACKAGE_ARC=$(echo "$SINGLE_PACKAGE" | cut -d':' -f2 | cut -d' ' -f1)
            PACKAGE_OLD_VER=$(echo "$SINGLE_PACKAGE" | cut -d'(' -f2 | cut -d',' -f1)
            PACKAGE_OLD_VER_WEB=$(echo "$PACKAGE_OLD_VER" | sed "s/^.*://") # Remove the "epoch" Prefix
            PACKAGE_OLD_VER_FILE=$(echo "$PACKAGE_OLD_VER" | sed "s/:/%3a/") # Replace the ':' "epoch" separator with the '%3a' html version
            PACKAGE_NEW_VER=$(echo "$SINGLE_PACKAGE" | cut -d',' -f2 | sed 's/^ //')
          }

          # Generate Log Text...
          {
            LOG_SINGLE_PACKAGE="$PACKAGE_NAME:$PACKAGE_ARC ($PACKAGE_NEW_VER, $PACKAGE_OLD_VER)"
            if [ "$LOG_COMMAND" == "" ]; then
              if [ "$COMMAND" == "Upgrade" ]; then
                LOG_COMMAND="Downgrade: $LOG_SINGLE_PACKAGE"
              else
                LOG_COMMAND="Upgrade: $LOG_SINGLE_PACKAGE"
              fi
            else
              LOG_COMMAND="$LOG_COMMAND, $LOG_SINGLE_PACKAGE"
            fi
          }

          # Check if the DEB file is stored in the APT System Cache...
          DEB_FILE=$(ls /var/cache/apt/archives/${PACKAGE_NAME}_${PACKAGE_OLD_VER_FILE}_*.deb 2>/dev/null)

          if [ "$DEB_FILE" == "" ]; then
            # Check if the DEB file is stored in the apt-rollback Packages Cache...
            DEB_FILE=$(ls $PKG_DIR/${PACKAGE_NAME}_${PACKAGE_OLD_VER_WEB}_*.deb 2>/dev/null)

            # Try to retrieve the DEB file from the Standard Repositories...
            if [ "$DEB_FILE" == "" ]; then
              echo -e "  $PACKAGE_NAME ver. $PACKAGE_OLD_VER is currently \e[1;31mNOT AVAILABLE\e[0m on the Local Caches"
              echo -n "  I'm looking in the Standard Repositories"
              apt-get install -y -qq --download-only --allow-downgrades $PACKAGE_NAME=$PACKAGE_OLD_VER &>/dev/null &
              Progress_Bar

              # Check if the DEB file is stored in the APT System Cache...
              DEB_FILE=$(ls /var/cache/apt/archives/${PACKAGE_NAME}_${PACKAGE_OLD_VER_FILE}_*.deb 2>/dev/null)

              # Try to retrieve the DEB file from Launchpad Web Site...
              if [ "$DEB_FILE" == "" ]; then
                echo -e " \e[1;31m[ NOT FOUND ]\e[0m"
                echo -n "  I'm looking on the Launchpad Web Site"
                PACK_REP_ULR=$(wget -q -O- "$REPOSITORY${PACKAGE_NAME}/$PACKAGE_OLD_VER" | grep "${PACKAGE_NAME}_${PACKAGE_OLD_VER_WEB}_" | grep -E "(all|$ARCHITECTURE).deb" | sed "s/^.*href/href/" | cut -d'"' -f2)
                PACKAGE_FILENAME=$(echo "$PACK_REP_ULR" | xargs basename 2>/dev/null)
                if [ ! "$PACKAGE_FILENAME" == "" ]; then
                  rm -f "/tmp/$PACKAGE_FILENAME"
                  wget -q "$PACK_REP_ULR" -O "/tmp/$PACKAGE_FILENAME" &
                  Progress_Bar
                  mv "/tmp/$PACKAGE_FILENAME" "$PKG_DIR" &>/dev/null
                  # Check if the DEB file is stored in the apt-rollback Packages Cache...
                  DEB_FILE=$(ls $PKG_DIR/${PACKAGE_NAME}_${PACKAGE_OLD_VER_WEB}_*.deb 2>/dev/null)
                fi

                if [ "$DEB_FILE" == "" ]; then
                  echo -e " \e[1;31m[ NOT FOUND ]\e[0m"
                  DPKG_COMMAND=""
                  break
                else
                  echo -e " \e[1;32m[ FOUND ]\e[0m"
                fi
              else
                echo -e " \e[1;32m[ FOUND ]\e[0m"
              fi
            fi
          fi

          if [ ! "$DPKG_COMMAND" == "" ]; then
            DPKG_COMMAND="$DPKG_COMMAND $DEB_FILE"
          else
            DPKG_COMMAND="$DEB_FILE"
          fi
        done

        if [ ! "$DPKG_COMMAND" == "" ]; then
          IFS="$DEFAULT_IFS"
          USER_NAME="$SUDO_USER"
          USER_ID=$(id -u "$USER_NAME")
          START_DATE=$(date "+%Y-%m-%d  %H:%M:%S")
          {
            IFS=" "
            dpkg -i $DPKG_COMMAND >/dev/null 2>"$TMP_FILE" &
            PID=$!
            IFS=$'\n'
            Progress_Bar
            wait $PID
            EXIT_CODE=$?
          }
          ERRORS_LOG=$(cat "$TMP_FILE" | grep -i -v "warning.*downgrad")
          rm -f "$TMP_FILE"
          if [ ! "$EXIT_CODE" == "0" ]; then
            echo "ERRORS:"
            echo "$ERRORS_LOG"
            exit 1
          else
            echo " done"
          fi
          if [ ! "$ERRORS_LOG" == "" ]; then
            echo
            echo "WARNINGS:"
            echo "$ERRORS_LOG"
          fi
          END_DATE=$(date "+%Y-%m-%d  %H:%M:%S")
          {
            echo
            echo "Start-Date: $START_DATE"
            echo "Commandline: dpkg -i $DPKG_COMMAND"
            echo "Requested-By: $USER_NAME ($USER_ID)"
            echo "$LOG_COMMAND"
            echo "End-Date: $END_DATE"
          } >> $LOG_FILE
        else
          echo
          echo -e "ERROR: Undo command \e[1;31mABORTED\e[0m"
          echo
          exit 1
        fi
        IFS=$'\n'
      fi
      IFS="$DEFAULT_IFS"
      echo -e "\e[0m"
      let "N=N+1"
    done
    MaintainCache
    echo "Successfully Undone!"
    echo
  fi
  exit
}

function Usage_Message()
{
  echo "Usage: apt-rollback [--last <n>] [--remove/--reinstall package-name] [--help]"
  echo
}

function Ask_Yes_No()
{
  WARNING=""
  while true; do
    read -r -p "$WARNING$1 [y/N]? " YESNO
    if [ -z "$YESNO" ]; then
      RESULT="n"
      break
    else
      case $YESNO in
        "y" | "Y")
          RESULT="y"
          break ;;
        "n" | "N")
          RESULT="n"
          break ;;
        *) WARNING="Please answer [y]es or [n]o."$'\n' ;;
      esac
    fi
  done
  echo
}

function First_Installation()
{
  if [ -f "$LOG_FILE.1.gz" ]; then
    COMMAND_GZIP=$(zgrep -m1 -E "Install: .*$1:" "$LOG_FILE.1.gz")"\n"
  else
    COMMAND_GZIP=""
  fi
  COMMAND_PLAIN=$(grep -m1 -E "Install: .*$1:" "$LOG_FILE")
  COMMAND=$(echo -e "$COMMAND_GZIP$COMMAND_PLAIN" | tail -1 | cut -d" " -f2- | sed "s/[(][^)]*[)]//g" | sed "s/ ,//g" | sed 's/ *$//g')
  echo "$COMMAND"
}

function Initial_Checks()
{
  if [ ! -f "$LOG_FILE" ]; then
    echo "The APT log file '$LOG_FILE' doesn't exist."
    echo
    exit 1
  fi
  if [ "$INSTALLATION_DATE" == "" ]; then 
    echo "WARNING: Cannot determinate the OS Installation Date"
    echo "         This will prevent to block UNDOing of Installation Packages"
    echo
  fi
}

function Show_Help()
{
  Usage_Message
  echo "  --last       Undo the last <n> APT commands"
  echo "               Supports the undo of the only Install, Remove and Purge commands"
  echo
  echo "  --remove     Remove an INSTALLED package and related configuration files"
  echo "               Removing also all its first installed dependencies"
  echo
  echo "  --reinstall  Reinstall a REMOVED package,"
  echo "               and all its first installed dependences"
  echo "               Reproducing exactly its first installation"
  echo
  echo "  --help       Print this help"
  echo
  exit
}

function Wrong_Parameters()
{
  echo "'$1' is a wrong parameter"
  echo
  Usage_Message
}
# --------------------------------------------------------------------------------

# Main Code...
# --------------------------------------------------------------------------------
echo "apt-rollback ver. $VERSION"
echo "Undo the last APT commands or a specified one"
echo

Initial_Checks

if [ $# -eq 0 ]; then
  echo "No arguments supplied."
  Usage_Message
  Ask_Yes_No "Do you wish to see the last APT command"
  if [ "$RESULT" == "y" ]; then
    Undo_Last_Command "1"
  fi
else
  case "$1" in
    "--last") Undo_Last_Command "$2" ;;
    "--remove") INSTALLED_PACKAGES=$(First_Installation "$2") ;;
    "--reinstall") REMOVED_PACKAGES=$(First_Installation "$2") ;;
    "--help") Show_Help ;;
    *) Wrong_Parameters "$1" ;;
  esac
fi

if [ -n "$INSTALLED_PACKAGES" ]; then
  echo -e "The selected APT command, performed the 'INSTALL' of the following packages: \e[1;32m$INSTALLED_PACKAGES\e[0m"
  echo
  Ask_Yes_No "Do you wish to REMOVE them?"
  if [ "$RESULT" == "y" ]; then
    echo -n -e "Working:\e[1;32m"
    # Remove last Installed Packages...
    apt-get purge -y -qq $INSTALLED_PACKAGES >/dev/null &
    Progress_Bar
    echo
    echo -e "\e[0mDone"
    echo
  fi
else
  if [ -n "$REMOVED_PACKAGES" ]; then
    echo -e "The selected APT command, performed the 'REMOVE' of the following packages: \e[1;32m$REMOVED_PACKAGES\e[0m"
    echo
    Ask_Yes_No "Do you wish to RE-INSTALL them?"
    if [ "$RESULT" == "y" ]; then
      echo -n -e "Working:\e[1;32m"
      # Install last Removed Packages...
      apt-get install -y -qq $REMOVED_PACKAGES >/dev/null &
      Progress_Bar
      echo
      echo -e "\e[0mDone"
      echo
    fi
  fi
fi
