#!/usr/bin/python
# -*- coding: utf-8 -*-
'''
Created on 2018. 10. 15.
@author: Hyechurn Jang, <hyjang@cisco.com>
'''

from __future__ import print_function
import os
import sys
import array
import fcntl
import socket
import struct
import docker
import argparse

PYG_BASE_DIR = '/opt/pygmalion'
PYG_TOKEN = PYG_BASE_DIR + '/token'
PYG_WORK_DIR = PYG_BASE_DIR + '/workspace'
PYG_INST_LOG = PYG_BASE_DIR + '/install.log'

def con_title(msg): print('%-32s' % msg, end=''); sys.stdout.flush()
def con_ok(): print(' --> [ OK ]'); sys.stdout.flush(); return True
def con_fail(): print(' --> [ Fail ]'); sys.stdout.flush(); return False

def bash(title, install_cmd, check_cmd=None, check_ret=0, output=PYG_INST_LOG, fail2exit=True):
    con_title(title)
    if os.system('%s >> %s 2>&1' % (install_cmd, output)): ret = con_fail()
    else:
        if check_cmd and check_ret != os.system('%s > /dev/null 2>&1' % (check_cmd)): ret = con_fail()
        else: ret = con_ok()
    if not ret and fail2exit: exit(ret)
    return ret

def prepare_env():
    print('PREPARE ENVIRONMENT')
    if os.path.exists(PYG_TOKEN):
        print('pygmalion already installed')
        exit(1)
    con_title('create directories')
    os.system('mkdir -p %s' % PYG_BASE_DIR)
    os.system('mkdir -p %s' % PYG_WORK_DIR)
    os.system('touch %s' % PYG_INST_LOG)
    con_ok()

def install_nfs_server():
    print('INSTALL NFS SERVER')
    bash('install nfs server',
         'apt install -y nfs-kernel-server',
         'dpkg -l | grep nfs-kernel-server')
    if os.system('cat /etc/exports | grep pygmalion > /dev/null 2>&1'):
        bash('register workspace',
             'echo "/opt/pygmalion/workspace    *(rw,sync,no_root_squash,no_subtree_check)"',
             'cat /etc/exports | grep pygmalion',
             output='/etc/exports')
    bash('restart nfs server',
         'systemctl restart nfs-kernel-server',
         'systemctl status -q --no-pager nfs-kernel-server')
    bash('register nfs service',
         'systemctl enable nfs-kernel-server')

def install_nfs_client(ip):
    print('INSTALL NFS CLIENT')
    bash('install nfs client',
         'apt install -y nfs-common',
         'dpkg -l | grep nfs-common')
    if os.system('cat /etc/fstab | grep pygmalion > /dev/null 2>&1'):
        bash('register workspace',
             'echo "%s:/opt/pygmalion/workspace /opt/pygmalion/workspace nfs defaults 0 0"' % ip,
             'cat /etc/fstab | grep pygmalion',
             output='/etc/fstab')
    bash('mount nfs server',
         'mount -a',
         'mount | grep pygmalion')

def install_packages():
    print('INSTALL PACKAGES')
    bash('install docker',
         'apt install -y docker.io',
         'dpkg -l | grep docker')
    bash('restart docker',
         'systemctl restart docker',
         'systemctl status -q --no-pager docker')
    bash('register docker service',
         'systemctl enable docker')
    if os.system('docker images | grep pygics > /dev/null 2>&1'):
        bash('download pygics image',
             'docker pull jzidea/pygics:latest',
             'docker images | grep pygics')

def get_ips():
    max_possible = 128
    bytes = max_possible * 32
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    names = array.array('B', '\0' * bytes)
    outbytes = struct.unpack('iL', fcntl.ioctl(s.fileno(), 0x8912, struct.pack('iL', bytes, names.buffer_info()[0])))[0]
    namestr = names.tostring()
    result = {}
    for i in range(0, outbytes, 40):
        name = namestr[i:i+16].split('\0', 1)[0]
        ip = namestr[i+20:i+24]
        ip = str(ord(ip[0])) + '.' + str(ord(ip[1])) + '.' + str(ord(ip[2])) + '.' + str(ord(ip[3]))
        result[ip] = name
    return result

def master_swarm(ip):
    os.system('docker swarm leave --force >> /dev/null 2>&1')
    bash('initialize docker swarm',
         """docker swarm init --advertise-addr %s | grep "\-\-token" | awk '{print $2}'""" % ip,
         'ls %s' % PYG_TOKEN,
         output=PYG_TOKEN)
    with open(PYG_TOKEN, 'r') as fd: token = fd.read()
    print('run command on worker node followed...')
    print('')
    print('    pmadm join -a %s -t %s' % (ip, token))
    print('')

def worker_swarm(ip, token):
    os.system('docker swarm leave --force >> /dev/null 2>&1')
    bash('initialize docker swarm',
         'docker swarm join --token %s %s:2377' % (token, ip))
    bash('save token',
         'echo "%s"' % token,
         'ls %s' % PYG_TOKEN,
         output=PYG_TOKEN)

def install_master(ip):
    ips = get_ips().keys()
    if ip not in ips:
        print('%s not in local IP addresses' % ip)
        print('local IP list : %s' % ', '.join(ips))
        exit(1)
    
    prepare_env()
    install_nfs_server()
    install_packages()
    master_swarm(ip)

def install_worker(ip, token):
    prepare_env()
    install_nfs_client(ip)
    install_packages()
    worker_swarm(ip, token)

def exit_node():
    if not os.path.exists(PYG_INST_LOG):
        print('pygmalion is not installed')
        exit(1)
    
    con_title('remove related system settings')
    os.system('docker swarm leave --force >> /dev/null 2>&1')
    os.system('sed -i "/pygmalion/d" /etc/fstab >> /dev/null 2>&1')
    os.system('sed -i "/pygmalion/d" /etc/exports >> /dev/null 2>&1')
    os.system('umount /opt/pygmalion/workspace >> /dev/null 2>&1')
    os.system('systemctl restart nfs-kernel-server >> /dev/null 2>&1')
    
    con_ok()
    bash('move token to old',
         'mv %s %s.old' % (PYG_TOKEN, PYG_TOKEN))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Pygmalion Admin Tool')
    parser.add_argument('command', help='[ init | join | exit ]')
    parser.add_argument('-a', '--address', default='', help='Master IP Address or Hostname')
    parser.add_argument('-t', '--token', default='', help='Token Generated on Master')
    args = parser.parse_args()
    cmd = args.command
    ip = args.address
    token = args.token
    
    if cmd == 'init' and ip != '': install_master(ip)
    elif cmd == 'join' and ip != '' and token != '': install_worker(ip, token)
    elif cmd == 'exit': exit_node()
    else: args.print_help()
