# -*- coding: utf-8 -*-
import StringIO
from fabric.api import *
from fabric.contrib import files
from application import db,app
from models.host_bastion_apply import HostBastionApply
from models.host_bastion_tasks import HostBastionTasks
from models.pubkey_changed_queue import PubkeyChangedQueue
from models.user_host import UserHost
from models.sshkey import Sshkey
from models.host import Host
from models.ip_address import IpAddress
from models.user import User
from views.functions import addmail
from views.pool import get_ansible_passwd
from views.host import get_hostinfo
from sqlalchemy import and_
from signal import signal, SIGTERM
import datetime,os,sys,atexit, getopt

ROOT_KEY_FILE = '/root/.ssh/authorized_keys'
EVANS_KEY_FILE = '/home/evans/.ssh/authorized_keys'
READONLY_KEY_FILE = '/home/readonly/.ssh/authorized_keys'
DELIMITER_BEGIN = '### Generated by SSH-KEY - BEGIN'
DELIMITER_END = DELIMITER_BEGIN.replace('BEGIN', 'END')
BASTION_PID_FILE = '/var/run/bastion_tasks.pid'
PUBKEY_PID_FILE = '/var/run/pubkey_changed.pid'
MAIL_PID_FILE = '/var/run/overdue_mail.pid'
env.user = 'root'
env.abort_on_prompts = False
#env.skip_bad_hosts = True
env.timeout = 30
env.debug = True
#env.password = '%s'%get_ansible_passwd().strip()
env.password = 'anjukeansible'
env.key_filename = '/root/.ssh/ansible_rsa'
env.abort_exception = True
env.disable_known_hosts = True

LANG = "en_US.UTF-8"



def atexit_removepid(pid_file):
    try:
        os.remove(pid_file)
    except:
        pass
    print '---------end--------------'

class Merge_Tasks_Pubkey:
    def run(self):
        print '---------start--------------'
        self.role_root_users = set()
        self.role_evans_users = set()
        self.role_readonly_users = set()
        param= None
        opts, args = getopt.getopt(sys.argv[1:], "hm:p:")
        for op, mod in opts:
            if op == "-p":
                param = mod

        if param == "bastion":
            self.bastion()
        elif param == "pubkey_update":
            self.pubkey_update()
        elif param == "overdue_mail":
            self.overdue_mail()

    def bastion(self):
        #阻止重复运行
        pid_file=BASTION_PID_FILE
        if os.path.isfile(pid_file):
            print("The program is running")
            return 1
        else:
          try:
            signal(SIGTERM, lambda signum, stack_frame: exit(1))
            atexit.register(lambda:atexit_removepid(pid_file))
            fd=os.open(pid_file,os.O_CREAT|os.O_EXCL|os.O_RDWR)
            os.write(fd,"%s\n" % os.getpid())
            os.close(fd)
          except:
            print("Cann't get a lock file")
            return 1


        now = datetime.datetime.now()
        task_status = HostBastionTasks.STATUS_FREE
        current_tasks = HostBastionTasks.query.filter(and_(HostBastionTasks.status == task_status, HostBastionTasks.exec_time <= now)).order_by(HostBastionTasks.id).limit(5).all()
        if current_tasks:
            for current_task in current_tasks:
                print '------Do Task:%s-----------'%current_task.id
                self.getTaskUids(current_task)
                flag = self.do_task(current_task.host_id)
                self.do_apply(current_task.apply_id,flag)
                if flag:
                    current_task.status = HostBastionTasks.STATUS_SUCCESS
                    if current_task.type == UserHost.STATUS_VALID:
                        userinfo = User.query.filter(User.id == current_task.uid).first()
                        subject = "【堡垒机权限】申请成功提示"
                        hostname = Host.query.filter(Host.id == current_task.host_id).first()
                        content = '您申请的主机：' + hostname.hostname + '的权限已开通成功'
                        addmail(userinfo.email, subject, content)
                    self.user_host_mapping(current_task.uid,current_task.host_id,current_task.role,current_task.type)
                else:
                    current_task.status = HostBastionTasks.STATUS_FAIL
                    email = app.config.get("MAIL_TO")['pubkey_fail']
                    subject = "【堡垒机权限】申请失败"
                    hostname = Host.query.filter(Host.id == current_task.host_id).first()
                    content = '您申请的主机：' + hostname.hostname + '的权限开通失败,请查明原因'
                    addmail(email, subject, content)
                db.session.commit()
                self.role_root_users = set()
                self.role_evans_users = set()
                self.role_readonly_users = set()


    def pubkey_update(self):
        #阻止重复运行
        pid_file=PUBKEY_PID_FILE
        if os.path.isfile(pid_file):
            print("The program is running")
            return 1
        else:
          try:
            signal(SIGTERM, lambda signum, stack_frame: exit(1))
            atexit.register(lambda:atexit_removepid(pid_file))
            fd=os.open(pid_file,os.O_CREAT|os.O_EXCL|os.O_RDWR)
            os.write(fd,"%s\n" % os.getpid())
            os.close(fd)
          except:
            print("Cann't get a lock file")
            return 1


        pubkey_queues = PubkeyChangedQueue.query.filter(PubkeyChangedQueue.my_host_flag == 0).limit(5).all()
        if pubkey_queues:
            uids = []
            for pubkey_queue in pubkey_queues:
                print '------Do Queue:%s-----------'%pubkey_queue.id
                tmp_uid = pubkey_queue.uid
                if tmp_uid not in uids:
                    uids.append(tmp_uid)
                pubkey_queue.my_host_flag = 1
                db.session.commit()
            if len(uids) <=0:
                return
            UserHosts = UserHost.query.filter(and_(UserHost.uid.in_(uids),UserHost.status == UserHost.STATUS_VALID)).all()
            if UserHosts:
                for userhost in UserHosts:
                    self.getPubkeyUids(userhost.host_id)
                    self.do_task(userhost.host_id)
                    self.role_root_users = set()
                    self.role_evans_users = set()
                    self.role_readonly_users = set()


    def do_task(self,host_id):
        try:
            params = dict()
            role_root_users = self.role_root_users
            role_evans_users = self.role_evans_users
            role_readonly_users = self.role_readonly_users
            #get ip
            host_info = Host.query.filter(Host.id == host_id).first()
            ip = self._get_ip(host_info.id)
            env.host = ip
            execute(self.check_remote_file, host=env.host)

            params['key_file'] = ROOT_KEY_FILE
            params['keys'] = self.get_keys(role_root_users)
            self.proceess_task(ip, params)

            #readonly
            check_res = execute(self.check_account_exist,'readonly',host = env.host)
            if check_res[env.host]:
                params['key_file'] = READONLY_KEY_FILE
                params['keys'] = self.get_keys(role_readonly_users)
                self.proceess_task(ip, params)

            check_res = execute(self.check_account_exist,'evans',host = env.host)
            if check_res[env.host]:
                evans_home = execute(self.process_check_evans_home, host=env.host)
                if evans_home[env.host]:
                    params['key_file'] = '/home/www/.ssh/authorized_keys'
                else:
                    params['key_file'] = EVANS_KEY_FILE

                params['keys'] = self.get_keys(role_evans_users)
                self.proceess_task(ip, params)
            return True
        except Exception , e:
            print e
            return False


    def _get_ip(self,host_id):
        host_info = Host.query.filter(Host.id == host_id).first()
        ip_info = IpAddress.query.filter(IpAddress.id == host_info.primary_ip_id).first()
        return ip_info.ipv4

    def process_check_evans_home(self):
            evans_home = run("cat /etc/passwd| grep evans")
            evans_list = evans_home.split(":")
            #if evans_list[5] == '/home/evans':
            if '/home/evans' in evans_home:
                return False
            return True


    def check_account_exist(self,account):
        flag = True
        with settings(hide('warnings', 'running', 'stderr'),warn_only=True):
            res = run("/usr/bin/id  %s"%account)
            if  'No such user' in res:
                flag = False
        return flag

    def get_keys(self,user_ids):
        keys = set()
        root_user_keys = Sshkey.query.filter(Sshkey.uid.in_(user_ids)).all()
        for root_user_key in root_user_keys:
                keys.add(root_user_key.key)
        return keys
    def proceess_task(self,host_ip, params):
        execute(self.build_host, params, host=env.host)

    def check_remote_file(self):
        self.create_remote_file(ROOT_KEY_FILE, 'root', 'root')
        self.create_remote_file(EVANS_KEY_FILE, 'evans', 'www-data')
        self.create_remote_file(READONLY_KEY_FILE, 'readonly', 'readonly')

    def create_remote_file(self,remote_file, user, group):
        remote_dir = remote_file[0:-16]
        if not files.exists(remote_dir, use_sudo=False, verbose=True):
            run("mkdir -p "+remote_dir+" && chown -R "+user+":"+group+" "+remote_dir+" && chmod 700 "+remote_dir)
        if not files.exists(remote_file, use_sudo=False, verbose=True):
            run("touch "+remote_file+" && chown "+user+":"+group+" "+remote_file+" && chmod 600 "+remote_file)

    def build_host(self,params):
        print 'Build host ip %s, %d keys.' % (env.host, len(params['keys']))
        f_in = StringIO.StringIO()
        get(params['key_file'], f_in)
        f_in.seek(0)  #defaults to 0 (absolute file positioning); other values are 1 (seek relative to the current position) and 2 (seek relative to the file's end).
        f_out = StringIO.StringIO()
        within = False
        for line in f_in:
            if within:
                if line.startswith(DELIMITER_END):
                    within = False

            elif line.startswith(DELIMITER_BEGIN):
                within = True

            else:
                f_out.write(line)
                if not line.endswith('\n'):
                    f_out.write('\n')

        if params['keys']:
            f_out.write(DELIMITER_BEGIN + '\n')
            for key in params['keys']:
                f_out.write(key + '\n' + '\n')
            f_out.write(DELIMITER_END + '\n')

        put(f_out, params['key_file'])

        f_in.close()
        f_out.close()

    def getTaskUids(self,current_task):
         if current_task.type == HostBastionTasks.TYPE_ADD:
            self.setUids(current_task.role,current_task.uid)
         else:
            self.user_host_mapping(current_task.uid,current_task.host_id,current_task.role,current_task.type)

         host_users = UserHost.query.filter(and_(UserHost.host_id == current_task.host_id, UserHost.status == UserHost.STATUS_VALID)).all()
         if host_users:
            for host_user in host_users:
                self.setUids(host_user.role,host_user.uid)

    def getPubkeyUids(self,host_id):
        host_users = UserHost.query.filter(and_(UserHost.host_id == host_id, UserHost.status == UserHost.STATUS_VALID)).all()
        if host_users:
            for host_user in host_users:
                self.setUids(host_user.role,host_user.uid)

    def setUids(self,role,uid):
        if int(role) == UserHost.ROLE_ROOT:
            self.role_root_users.add(uid)
            self.role_evans_users.add(uid)
            self.role_readonly_users.add(uid)

        elif int(role) == UserHost.ROLE_EVANS:
            self.role_evans_users.add(uid)
            self.role_readonly_users.add(uid)

        elif int(role) == UserHost.ROLE_READONLY:
            self.role_readonly_users.add(uid)


    def user_host_mapping(self,uid,host_id,role,status):
        hasIn = UserHost.query.filter(and_(UserHost.uid==uid,UserHost.host_id==host_id,UserHost.role==role)).first()
        if status == UserHost.STATUS_VALID:#新增
            if not hasIn:
                target = UserHost(uid,host_id,role,status)
                db.session.add(target)
                db.session.commit()
            else:
                hasIn.status = UserHost.STATUS_VALID
                db.session.commit()

        else:#删除
            if hasIn:
                hasIn.status = UserHost.STATUS_OVERDUE
                db.session.commit()

    def do_apply(self,apply_id,flag):
        hasIn = HostBastionApply.query.filter(HostBastionApply.id == apply_id).first()
        if not hasIn:
            return
        if flag:
            hasIn.status = HostBastionApply.STATUS_SUCCESS
        else:
            hasIn.status = HostBastionApply.STATUS_FAIL
        db.session.commit()

    def overdue_mail(self):
        #阻止重复运行
        pid_file=MAIL_PID_FILE
        if os.path.isfile(pid_file):
            print("The program is running")
            return 1
        else:
          try:
            signal(SIGTERM, lambda signum, stack_frame: exit(1))
            atexit.register(lambda:atexit_removepid(pid_file))
            fd=os.open(pid_file,os.O_CREAT|os.O_EXCL|os.O_RDWR)
            os.write(fd,"%s\n" % os.getpid())
            os.close(fd)
          except:
            print("Cann't get a lock file")
            return 1

        today = datetime.date.today()
        exctime = today + datetime.timedelta(3)
        begin_time = str(exctime) + " 00:00:00"
        end_time = str(exctime) + " 23:59:59"
        print '-------Task  in (%s,%s)-----------'%(begin_time,end_time)
        task_status = HostBastionTasks.STATUS_FREE
        task_type = HostBastionTasks.TYPE_DELETE
        pre_exc_tasks = HostBastionTasks.query.filter(and_(HostBastionTasks.status == task_status, HostBastionTasks.type == task_type, HostBastionTasks.exec_time.between(begin_time,end_time))).order_by(HostBastionTasks.id).limit(5).all()
        uids = []
        userhost = {}
        if pre_exc_tasks:
            for item in pre_exc_tasks:
                uids.append(item.uid)
                hostinfo = get_hostinfo(item.host_id)
                if item.uid not in userhost:
                    userhost[item.uid] = {'host_id':[],'hostname':[]}
                userhost[item.uid]['host_id'].append(item.host_id)
                userhost[item.uid]['hostname'].append(hostinfo.hostname)

            userinfo = User.query.filter(User.id.in_(uids)).all()
            if userinfo:
                print '---------开始发邮件-------------'
                DOMAIN = app.config.get("DOMAIN")
                for item in userinfo:
                    receiver = item.email.replace(',',';')
                    subject = "【堡垒机权限过期】3天后权限到期主机列表提醒"
                    names = ','.join(userhost[item.id]['hostname'])
                    content = "您申请的主机"+names+"的权限还有3天过期,请及时续期限,保证权限正常使用"
                    addmail(receiver, subject, content)
                print('-------邮件发送结束------')
