from email.mime.text import MIMEText

import argparse
import os
import re
import smtplib
import subprocess
import sys
import threading
import time


# default values
NFS_LOCATION = "/mnt/ha"
HEAD1 = "10.0.0.21"
HEAD2 = "10.0.0.22"
HEAD1_IPMI = "10.0.0.23"
HEAD2_IPMI = "10.0.0.24"
HEAD1_IPMIUSER = 'ADMIN'
HEAD2_IPMIUSER = 'ADMIN'
HEAD1_IPMIPW = "abcd1234"
HEAD2_IPMIPW = "abcd1234"

EMAIL_HOST = 'localhost'
EMAIL_SSL = False
EMAIL_FROM = 'me@localhost'
EMAIL_TO = 'me@localhost'
EMAIL_USER = 'me@localhost'
EMAIL_PASSWORD = 'me@localhost'


class NFSMount(threading.Thread):

    def __init__(self, location):
        self._location = location
        self.retval = None
        super(NFSMount, self).__init__()

    def run(self):

        _file = "%s/.test" % self._location
        with open(_file, 'w') as f:
            f.write('test\n')

        os.unlink(_file)

        self.retval = True


class MasterHead(threading.Thread):

    def __init__(self, host):
        self._host = host
        self._proc = None
        self.retval = None
        super(MasterHead, self).__init__()

    def run(self):

        self._proc = subprocess.Popen(["ssh",
            self._host,
            "ifconfig",
            "carp0",
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)

        output, err = self._proc.communicate()
        if self._proc.returncode != 0:
            self.retval = False
            return

        reg = re.search(r'\bcarp: (?P<status>\w+)\s', output, re.S)
        if not reg:
            self.retval = False
            return

        status = reg.group("status")
        self.retval = status

    def finish(self):
        if self._proc and self.retval is None:
            try:
                os.kill(self._proc.pid, 9)
            except OSError:
                pass


class PowerReset(threading.Thread):

    def __init__(self, ip, user, password):
        self._ip = ip
        self._password = password
        self._user = user
        self.retval = None
        super(PowerReset, self).__init__()

    def run(self):

        self._proc = subprocess.Popen(["ipmitool",
            "-H", self._ip,
            "-U", self._user,
            "-P", self._password,
            "chassis",
            "power",
            "reset",
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
        if self._proc.wait() == 0:
            self.retval = True
        else:
            self.retval = False


def reboot(ip, user, password):

    pr = PowerReset(ip, user, password)
    pr.start()
    pr.join(5)
    if pr.retval is None:
        try:
            os.kill(pr._proc.pid, 9)
        except OSError:
            pass
    return (pr.retval is True)


def sendmail(msgs):
    msg = MIMEText('\n'.join(msgs))
    msg['Subject'] = 'TrueNAS HA STONITH'
    msg['From'] = EMAIL_FROM
    msg['To'] = EMAIL_TO
    if EMAIL_SSL:
        server = smtplib.SMTP_SSL(EMAIL_HOST)
    else:
        server = smtplib.SMTP(EMAIL_HOST)
    if EMAIL_USER:
        server.login(EMAIL_USER, EMAIL_PASSWORD)
    #server.set_debuglevel(1)
    server.sendmail(EMAIL_FROM, [EMAIL_TO], msg.as_string())
    server.quit()

def main():

    parser = argparse.ArgumentParser()
    parser.add_argument("--nfs", "-n", action='store', default=NFS_LOCATION)
    parser.add_argument("--head1", "-h1", action='store', default=HEAD1)
    parser.add_argument("--head2", "-h2", action='store', default=HEAD2)
    parser.add_argument("--head1ipmi", "-h1i", action='store', default=HEAD1_IPMI)
    parser.add_argument("--head2ipmi", "-h2i", action='store', default=HEAD2_IPMI)
    parser.add_argument("--head1ipmiuser", "-h1iuser", action='store', default=HEAD1_IPMIUSER)
    parser.add_argument("--head2ipmiuser", "-h2iuser", action='store', default=HEAD2_IPMIUSER)
    parser.add_argument("--head1ipmipw", "-h1ipw", action='store', default=HEAD1_IPMIPW)
    parser.add_argument("--head2ipmipw", "-h2ipw", action='store', default=HEAD2_IPMIPW)
    args = parser.parse_args()
    print args

    logs = []
    nfsmount = NFSMount(args.nfs)
    nfsmount.start()

    # Wait for 5 seconds
    nfsmount.join(5)

    if nfsmount.retval is True:
        sys.exit(0)

    logs.append("NFS Mount hung: %s" % args.nfs)

    head1 = MasterHead(args.head1)
    head1.start()

    head2 = MasterHead(args.head2)
    head2.start()

    head1.join(5)
    head1.finish()
    head2.join(5)
    head2.finish()

    reboot_head1 = False
    reboot_head2 = False

    # None means process hung
    if head1.retval in (None, 'MASTER'):
        reboot_head1 = True

    # None means process hung
    if head2.retval in (None, 'MASTER'):
        reboot_head2 = True

    if reboot_head1:
        logs.append("head1 rebooting")
        reboot(args.head1ipmi, args.head1ipmiuser, args.head1ipmipw)
        time.sleep(10)
        head2 = MasterHead(args.head1)
        head2.start()
        head2.join(5)
        head2.finish()
        if head2.retval != 'MASTER':
            logs.append("head2 didnt take over, carp status %s" % head2.retval)
        else:
            logs.append("head2 take over")

    if reboot_head2:
        logs.append("head2 rebooting")
        reboot(args.head2ipmi, args.head2ipmiuser, args.head2ipmipw)
        time.sleep(10)
        head1 = MasterHead(args.head1)
        head1.start()
        head1.join(5)
        head1.finish()
        if head1.retval != 'MASTER':
            logs.append("head1 didnt take over, carp status %s" % head1.retval)
        else:
            logs.append("head1 take over")

    if logs:
        print logs
        sendmail(logs)

if __name__ == "__main__":
    main()
