#!/usr/bin/python3 -u

# Changes the boot order to revert to the previous OSTree deployment.

# On Endless OS, after the first OS upgrade, we keep two deployments available.
# So, each call to this script will toggle between the two deployments.

from argparse import ArgumentParser
import gi
gi.require_version('OSTree', '1.0')
from gi.repository import Gio, OSTree
import os
import sys
import time

def print_deployments(deployments):
    for index, deploy in enumerate(deployments):
        print('Deployment {}:'.format(index))
        print(' Index:', deploy.get_index())
        print(' Checksum:', deploy.get_csum())
        print(' Boot checksum:', deploy.get_bootcsum())
        print(' OS:', deploy.get_osname())
        print(' Deploy serial:', deploy.get_deployserial())
        print(' Boot serial:', deploy.get_bootserial())

aparser = ArgumentParser(
    description='Rollback to previous OSTree deployment'
)
aparser.add_argument('--sysroot', help='path to OSTree sysroot')
args = aparser.parse_args()

if os.getuid() != 0:
    print('Program requires superuser privileges')
    sys.exit(1)

print()
response = input('Are you sure you want to rollback to the previous OS version? [y/N] ')
response = response.lower()
if response != 'yes' and response != 'y':
    sys.exit(1)

sysroot_file = None
if args.sysroot is not None:
    sysroot_file = Gio.File.new_for_path(args.sysroot)
sysroot = OSTree.Sysroot.new(sysroot_file)
sysroot.load()
sysroot_path = sysroot.get_path().get_path()

# Lock the sysroot, waiting up to 5 minutes
wait = 5 * 60
while True:
    locked = sysroot.try_lock()
    if locked:
        break
    if wait <= 0:
        print('Could not lock OSTree sysroot', sysroot_path,
              file=sys.stderr)
        sys.exit(1)

    # Try again in 1 second
    if wait % 60 == 0:
        minutes = wait / 60
        print('Waiting', minutes, 'minute' + 's' if minutes > 1 else '',
              'to acquire OSTree sysroot lock for', sysroot_path)
    wait -= 1
    time.sleep(1)

deployments = sysroot.get_deployments()
print()
print('Old deployments...')
print_deployments(deployments)
if len(deployments) == 0:
    print('No OSTree deployments found', file=sys.stderr)
    sys.exit(1)
elif len(deployments) == 1:
    print('Only one OSTree deployment found, cannot rollback',
          file=sys.stderr)
    sys.exit(1)

# Construct the new deployment list with the previous deployment first,
# rotating so that the current deployment is at the end of the list
new_deployments = deployments[1:] + [deployments[0]]
print()
print('New deployments...')
print_deployments(new_deployments)

# Now write the new deployments
print()
sysroot.write_deployments(new_deployments)

print()
print('Rollback complete. Please reboot the computer to start the previous OS version.')
