#!/usr/bin/python
#
# 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("/etc/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", "/etc/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", "/etc/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 /etc/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", "/etc/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()
