#!/usr/bin/python3
#
# Copyright (c) 2017-2019, Virtuozzo International GmbH
#
# Our contact details: Virtuozzo International GmbH, Vordergasse 59, 8200
# Schaffhausen, Switzerland.
#

import argparse
import os
import shutil
import sys

def vz_release_file():
    if os.path.isfile("/etc/vz-platform-release") or os.path.isfile("/etc/vz-chameleon/vz-platform-release"):
        return "vz-platform-release"
    else:
        return "vzlinux-release"

def backup_vzlinux():
    shutil.move("/etc/" + vz_release_file(), "/etc/vz-chameleon/" + vz_release_file())
    shutil.move("/usr/lib/os-release", "/etc/vz-chameleon/vzlinux-os-release")
    shutil.move("/etc/redhat-release", "/etc/vz-chameleon/vzlinux-redhat-release")

def mimic_rhel():
    if not os.path.isfile("/etc/vz-chameleon/chameleon-redhat-release"):
        print("/etc/vz-chameleon/chameleon-redhat-release is missing, cannot mimic RHEL")
        sys.exit(1)
    if not os.path.isfile("/etc/vz-chameleon/chameleon-redhat-os-release"):
        print("/etc/vz-chameleon/chameleon-redhat-os-release is missing, cannot mimic RHEL")
        sys.exit(1)
    if os.path.isfile("/etc/" + vz_release_file()):
        backup_vzlinux()
    shutil.copy2("/etc/vz-chameleon/chameleon-redhat-release", "/etc/redhat-release")
    shutil.copy2("/etc/vz-chameleon/chameleon-redhat-os-release", "/usr/lib/os-release")
    if os.path.isfile("/etc/centos-release"):
        os.remove("/etc/centos-release")
    print("RHEL chameleon mode is enabled")

def mimic_centos():
    if not os.path.isfile("/etc/vz-chameleon/chameleon-centos-release"):
        print("/etc/vz-chameleon/chameleon-centos-release is missing, cannot mimic CentOS")
        sys.exit(1)
    if not os.path.isfile("/etc/vz-chameleon/chameleon-centos-os-release"):
        print("/etc/vz-chameleon/chameleon-centos-os-release is missing, cannot mimic CentOS")
        sys.exit(1)
    if os.path.isfile("/etc/" + vz_release_file()):
        backup_vzlinux()
    shutil.copy2("/etc/vz-chameleon/chameleon-centos-release", "/etc/centos-release")
    shutil.copy2("/etc/vz-chameleon/chameleon-centos-release", "/etc/redhat-release")
    shutil.copy2("/etc/vz-chameleon/chameleon-centos-os-release", "/usr/lib/os-release")
    print("CentOS chameleon mode is enabled")

def disable_mimic():
    if os.path.isfile("/etc/" + vz_release_file()):
        print("Chameleon mode is already disabled, nothing to do")
        sys.exit(1)
    if not os.path.isfile("/etc/vz-chameleon/" + vz_release_file()):
        print("/etc/vz-chameleon/" + vz_release_file() + " is missing, cannot restore /etc/" + vz_release_file())
        sys.exit(1)
    if not os.path.isfile("/etc/vz-chameleon/vzlinux-os-release"):
        print("/etc/vz-chameleon/vzlinux-os-release is missing, cannot restore /usr/lib/os-release")
        sys.exit(1)
    if not os.path.isfile("/etc/vz-chameleon/vzlinux-redhat-release"):
        print("/etc/vz-chameleon/vzlinux-redhat-release is missing, cannot restore /etc/redhat-release")
        sys.exit(1)
    if os.path.isfile("/etc/centos-release"):
        os.remove("/etc/centos-release")
    shutil.move("/etc/vz-chameleon/" + vz_release_file(), "/etc/" + vz_release_file())
    shutil.move("/etc/vz-chameleon/vzlinux-os-release", "/usr/lib/os-release")
    shutil.move("/etc/vz-chameleon/vzlinux-redhat-release", "/etc/redhat-release")
    print("Chameleon mode is disabled")

if __name__ == '__main__':
    if os.geteuid() != 0:
        print("You need to have root privileges to run this script!")
        sys.exit(1)

    parser = argparse.ArgumentParser(description="Helper tool to make distribution mimic RHEL or CentOS.")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--enable', action='store', choices=['rhel', 'centos'], help='Mimic RHEL or CentOS')
    group.add_argument('--disable', action='store_true', help='Do no mimic RHEL or CentOS')
    args = parser.parse_args(sys.argv[1:])

    if args.disable:
        disable_mimic()

    if args.enable == 'rhel':
        mimic_rhel()

    if args.enable == 'centos':
        mimic_centos()
