#!/usr/libexec/platform-python
import argparse
import sys
import os
import re
import requests
from datetime import datetime, timedelta
import dnf


def getOSInfo():
  # Get the dnf variable name and symlink prefix for the current OS
  with open('/etc/os-release') as f:
    data = {k:v.strip('"\n') for k,v in (l.split('=') for l in f if l.strip())}

  if data['NAME'] == 'AlmaLinux':
    return ('cernalmalinux', data['VERSION_ID'].split('.')[0])
  elif data['NAME'] == 'Red Hat Enterprise Linux':
    return ('cernrhel', data['VERSION_ID'].split('.')[0])
  elif data['NAME'] == 'CentOS Stream':
    if data['VERSION'] == '8':
      return ('cernstream8', 's8')
    # Must be Stream 9
    return ('cernstream9', 's9')
  else:
    print("Sorry: %s is not supported" % data['NAME'])
    sys.exit(1)

def getBaseURL(variable = None, value = None):
  base = dnf.Base()
  base.conf.substitutions.update_from_etc('/')
  if variable and value:
    base.conf.substitutions[variable] = value
  base.read_all_repos()
  base.fill_sack()
  try:
    repo = base.repos.get('baseos')
    return repo.baseurl[0]
  except:
    print("Unable to find the baseos repository.")
    sys.exit(1)

def dnfVar(method, variable, value=None):
  if not os.path.isfile('/etc/dnf/vars/%s' % variable):
    print('Unable to open /etc/dnf/vars/%s. Is the correct *-release RPM installed from the CERN repository?' % variable)
    sys.exit(1)

  if method == 'r':
    print("Your system is currently configured to use %s" % getBaseURL())
  elif method == 'r+':
    rootCheck()
    checkRepo(variable, value)
    with open('/etc/dnf/vars/%s' % variable, method) as dnfvar:
      dnfvar.truncate(0)
      dnfvar.write(value)
  else:
    print('%s is not a supported method' % method)
    sys.exit(1)

def rootCheck():
  if os.geteuid() != 0:
    print("Please run this script as root or sudo to modify the dnf variable")
    sys.exit(1)

def checkRepo(variable, value):
    url = getBaseURL(variable, value)
    response = requests.get(url)
    if response.status_code != 200:
      print("Sorry: %s was not found" % (url))
      sys.exit(1)

def main():
  parser = argparse.ArgumentParser(description='Show and modify .repo files to allow staged updates for CERN')
  parser.add_argument('-s', '--show', required=False, action='store_true', help='show current config')
  parser.add_argument('-p', '--production', required=False,action='store_true', help='configure production')
  parser.add_argument('-t', '--testing', required=False,action='store_true', help='configure testing')
  parser.add_argument('-d', '--days', required=False,type=str,action='store', help='configure custom delay based on x days or YYYYMMDD')
  parser.add_argument('-r', '--release', required=False,type=str,action='store', help='configure repositories to point to specific minor release; eg 8.0, 8.1, 8.2, etc')
  parser.add_argument('--prerelease', required=False,action='store_true', help=argparse.SUPPRESS)
  args = parser.parse_args()

  OSINFO = getOSInfo()

  if args.show:
    dnfVar('r', OSINFO[0])
    sys.exit(0)
  elif args.production:
    dnfVar('r+', OSINFO[0], OSINFO[1])
  elif args.testing:
    dnfVar('r+', OSINFO[0], '%s-testing' % OSINFO[1])
  elif args.days:
    # date
    if re.search(r'^\d{8}$', args.days) is not None:
      snapshot_path = '%s-snapshots/%s' % (OSINFO[1], args.days)
      dnfVar('r+', OSINFO[0], snapshot_path)
    # days delta
    elif re.search(r'^\d*$', args.days) is not None:
      d = datetime.today() - timedelta(days=int(args.days))
      snapshot_path = '%s-snapshots/%s' % (OSINFO[1], d.strftime('%Y%m%d'))
      dnfVar('r+', OSINFO[0], snapshot_path)
    else:
      print("Sorry: Please provide a date in the format of YYYYMMDD or an offset in days")
      sys.exit(1)
  elif args.release:
    if OSINFO[0] not in ['cernrhel', 'cernalmalinux']:
      print("Sorry, releases are only available for RHEL or AlmaLinux")
      sys.exit(1)
    # 8.1 or 8.1.1911
    if re.search(r'^\d.\d+$|^\d.\d.\d+$', args.release) is None:
      print("Sorry: A minor release should be passed as X.$release: example '8.0'")
      sys.exit(1)
    dnfVar('r+', OSINFO[0], args.release)
    print("Note: Setting the dnf variable to a specific release is not officially supported, and is only provided as a convenience")
    print("Note: Updates will not be applied when configured against a specific release")
  elif args.prerelease:
    print("Note: Setting the dnf variable to prerelease is NOT supported. I hope you know what you are doing ...")
    dnfVar('r+', OSINFO[0], '.%s-latest' % OSINFO[1])

  else:
    parser.print_help(sys.stdout)
  sys.exit(0)

if __name__ == "__main__":
    main()
