#!/usr/bin/python3

###########################################################
#
#              Setup Debian mirrors
#
# By: VIGIL NVR Team
#
# Dependency: elevate
#                maybe requires: sudo pip3 install elevate
#
############################################################

import os, sys, hashlib
import elevate

TEST = False

try:
  from elevate import elevate
except:
  print("You need to install \"elevate\" package before running this script, run:")
  print("  sudo pip3 install elevate")
  print("then rerun this script.")

# elevate package is required:  sudo pip3 install elevate

if TEST == True:
  SECURITY_LIST_PATH = "test-sec.list"
else:
  SECURITY_LIST_PATH = "/etc/apt/sources.list.d/debian-security.list"

SECURITY_LIST_CONTENT = """
deb http://security.debian.org/debian-security/ bullseye-security main contrib non-free
"""

if TEST == True:
  DEBIAN_PUBLIC_PATH = "test.list"
else:
  DEBIAN_PUBLIC_PATH = "/etc/apt/sources.list"

DEBIAN_PUBLIC_CONTENT = """
deb http://deb.debian.org/debian bullseye main contrib non-free
deb http://deb.debian.org/debian bullseye-updates main contrib non-free
deb http://deb.debian.org/debian bullseye-backports main contrib non-free
"""

MSG_ROOT_WARNING = """
You need to run this script as a root. Use sudo command:
  sudo python3 update_mirror.py
"""

MSG_ADD_SECURITY = "Do you want to enable the Debian public security mirror?"
MSG_SWICTH_TO_PUBLIC = "Do you want to switch to all Debian public mirrors?"

MSG_WARNING = """

**************************************************************
  !!! These changes are irreversible. Future updates may
  require monitoring and manual intervention. In some cases,
  your NVR/DVR may become non-operational if Debian packages
  are broken or if the installer prompts for user input to
  select system options. While such situations are rare, 
  they are beyond 3xLOGIC's control.
**************************************************************

Do you want to proceed ?"""

MSG_SECOND_WARNING = "Second confirmation: do you really want to proceed?"


#shutil.copyfile("debian-security.list", "/etc/apt/sources.list.d/debian-security.list")
DISTRO_ID_FILENAME = "/etc/issue"
DISTRO_DEBIAN_11_PATTERN = "Debian GNU/Linux 11"

def check_distro():
    f = open(DISTRO_ID_FILENAME)
    d = f.read()
    f.close()
    idx = d.find(DISTRO_DEBIAN_11_PATTERN)
    if idx < 0:
        print("This script should only be used under Debian 11")
        sys.exit(1)

def get_text_md5(text):
  return hashlib.md5(text.encode('utf-8')).hexdigest()

def get_file_md5(filepath):
  try:
    with open(filepath, "rb") as f:
        file_bytes = f.read()
        md5_hash = hashlib.md5(file_bytes).hexdigest()
    return md5_hash
  except:
    pass
  return "0"


def update():
  if TEST == True:
    print("Test: it would updating the system")
  else:
    print("Updating the system")
    os.system("apt update")
    os.system("apt upgrade -y")    

def yes_true_or_no_false(msg):
  user_input = input("\n" + msg + " (yes/no): ")
  if user_input.lower() in ["yes", "y"]:
      return True
  return False

def set_security_mirror():
  try:
    f = open(SECURITY_LIST_PATH, "w")
    f.write(SECURITY_LIST_CONTENT)
    f.close()
  except:
    print("Fail to write to " + SECURITY_LIST_PATH)
    sys.exit(1)

def set_all_public_mirrors():
  try:
    f = open(DEBIAN_PUBLIC_PATH, "w")
    f.write(DEBIAN_PUBLIC_CONTENT)
    f.close()
  except:
    print("Fail to write to " + DEBIAN_PUBLIC_PATH)
    sys.exit(1)
  
def main():
  add_security = False
  add_all_public = False
  # check if security repo already set
  if os.path.exists(SECURITY_LIST_PATH) == False:
     add_security = yes_true_or_no_false(MSG_ADD_SECURITY)
  else:
     print("\nThis system is already using Debian public security mirror.")
  
  # check if we already have setup all public mirrors
  if get_file_md5(DEBIAN_PUBLIC_PATH) != get_text_md5(DEBIAN_PUBLIC_CONTENT):
    add_all_public = yes_true_or_no_false(MSG_SWICTH_TO_PUBLIC)
  else:
    print("This system is already using All Debian public mirrors.")
  
  if add_all_public == False and add_security == False:
    print("Nothing to be done")
    return
  
  if add_all_public == True:
    if os.path.exists(SECURITY_LIST_PATH) == False:
      print("Security mirror will also be setup with all public mirrors")
      add_security = True
  
  if yes_true_or_no_false(MSG_WARNING) == False:
    print("Nothing has been changed, exiting")
    sys.exit(0)

  if yes_true_or_no_false(MSG_SECOND_WARNING) == False:
    print("Nothing has been changed, exiting")
    sys.exit(0)
  
  if add_all_public == True:
    set_all_public_mirrors()

  if add_security == True:
    set_security_mirror()
  
  update()
  print("All done")
  

if __name__ == "__main__":
  try:
    elevate(graphical=False)
  except:
    print(MSG_ROOT_WARNING)
    sys.exit(1)

  if os.geteuid() != 0:
    print(MSG_ROOT_WARNING)
    sys.exit(1)

  check_distro()
  main()