#!/usr/bin/python2.7

# Copyright (c) 2014, 2015, Oracle and/or its affiliates. All rights reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import errno
import os
import pwd
import smf_include
import subprocess
import sys
import time

from subprocess import CalledProcessError, check_call, Popen
from tempfile import mkstemp

GTF = "/usr/bin/gtf"
SVCCFG = "/usr/sbin/svccfg"
SVCPROP = "/usr/bin/svcprop"
VNCSERVER = "/usr/bin/vncserver"
XRANDR = "/usr/bin/xrandr"
XSTARTUPHDR = "# WARNING: THIS FILE GENERATED BY SMF.\n" + \
              "#   DO NOT EDIT THIS FILE.  EDITS WILL BE LOST.\n"
XRESOURCES = "[[ -f ~/.Xresources ]] && /usr/bin/xrdb -merge ~/.Xresources\n"
XTERM = "/usr/bin/xterm"
# Borderless, Monospsce font, point size 14, white foreground on black
# background are reasonable defaults.
XTERMOPTS = ' -b 0 -fa Monospace -fs 14 -fg white -bg black -title ' + \
            '"Zone Console: $ZONENAME"'
XWININFO = "/usr/bin/xwininfo"
# Enclose command in comments to prevent xterm consuming zlogin opts
ZLOGINOPTS = ' -e "/usr/bin/pfexec /usr/sbin/zlogin -C -E $ZONENAME"\n'
XSTARTUP = XSTARTUPHDR + XRESOURCES + XTERM + XTERMOPTS + ZLOGINOPTS


def start():
    check_vncserver()
    homedir = os.environ.get('HOME')
    if not homedir:
        homedir = pwd.getpwuid(os.getuid()).pw_dir
        os.putenv("HOME", homedir)
    set_xstartup(homedir)

    try:
        fmri = os.environ['SMF_FMRI']
        zonename = fmri.rsplit(':', 1)[1]
        os.putenv("ZONENAME", zonename)
        desktop_name = zonename + ' console'
        # NOTE: 'geometry' below is that which matches the size of standard
        # 80 character undecorated xterm window using font style specified in
        # XTERMOPTS. The geometry doesn't matter too much because the display
        # will be resized using xrandr once the xterm geometry is established.
        cmd = [VNCSERVER, "-name", desktop_name, "-SecurityTypes=None",
               "-geometry", "964x580", "-localhost", "-autokill"]
        vnc = Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                    env=None)
        out, err = vnc.communicate()
        vncret = vnc.wait()
        if vncret != 0:
            print "Error starting VNC server: " + err
            return smf_include.SMF_EXIT_ERR_FATAL
    except Exception as e:
        print e
        return smf_include.SMF_EXIT_ERR_FATAL

    output = err.splitlines()
    for line in output:
        if line.startswith("New '%s' desktop is" % desktop_name):
            display = line.rpartition(' ')[2]
            host, display_num = display.split(':', 1)
            # set host prop
            port = 5900 + int(display_num)
            print "VNC port: %d" % port
            # set port num prop
            cmd = [SVCCFG, '-s', fmri, 'setprop', 'vnc/port', '=', 'integer:',
                   str(port)]

            svccfg = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                      stderr=subprocess.PIPE)
            out, err = svccfg.communicate()
            retcode = svccfg.wait()
            if retcode != 0:
                print "Error updating 'vnc/port' property: " + err
                return smf_include.SMF_EXIT_ERR_FATAL
    resize_xserver(display, zonename)

    return smf_include.SMF_EXIT_OK


def stop():
    try:
        # first kill the SMF contract
        check_call(["/usr/bin/pkill", "-c", sys.argv[2]])
    except CalledProcessError as cpe:
        # 1 is returncode if no SMF contract processes were matched,
        # meaning they have already terminated.
        if cpe.returncode != 1:
                print "failed to kill the SMF contract: %s" % cpe
                return smf_include.SMF_EXIT_ERR_FATAL

    try:
        fmri = os.environ['SMF_FMRI']
        # reset port num prop to initial zero value
        cmd = [SVCCFG, '-s', fmri, 'setprop', 'vnc/port', '=', 'integer:',
               '0']
        svccfg = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                  stderr=subprocess.PIPE,)
        out, err = svccfg.communicate()
        retcode = svccfg.wait()
        if retcode != 0:
            print "Error resetting 'vnc/port' property: " + err
            return smf_include.SMF_EXIT_ERR_FATAL
    except Exception as e:
        print e
        return smf_include.SMF_EXIT_ERR_FATAL


def check_vncserver():
    if not os.path.exists(VNCSERVER):
        print("VNC console service not available on this compute node. "
              "%s is missing. Run 'pkg install x11/server/xvnc'"
              % VNCSERVER)
        return smf_include.SMF_EXIT_ERR_FATAL
    if not os.path.exists(XTERM):
        print("VNC console service not available on this compute node. "
              "%s is missing. Run 'pkg install terminal/xterm'"
              % XTERM)
        return smf_include.SMF_EXIT_ERR_FATAL


def set_xstartup(homedir):
    vncdir = os.path.join(homedir, '.vnc')
    xstartup_path = os.path.join(vncdir, 'xstartup')

    try:
        os.mkdir(vncdir)
    except OSError as ose:
        if ose.errno != errno.EEXIST:
            raise

    # Always clobber xstartup
    # stemp tuple = [fd, path]
    stemp = mkstemp(dir=vncdir)
    os.write(stemp[0], XSTARTUP)
    os.close(stemp[0])
    os.chmod(stemp[1], 0700)
    os.rename(stemp[1], xstartup_path)


def resize_xserver(display, zonename):
    """ Try to determine xterm window geometry and resize the Xvnc display
        to match using XRANDR. Treat failure as non-fatal since an
        incorrectly sized console is arguably better than none.
    """
    class UninitializedWindowError(Exception):
        pass

    class UnmappedWindowError(Exception):
        pass

    def _get_window_geometry(display, windowname):
        """ Find the xterm xwindow by name/title and extract its geometry
            Returns: tuple of window [width, height]
            Raises:
                UninitializedWindowError if window not yet initialized
                UnmappedWindowError if window is not viewable/mapped
        """
        cmd = [XWININFO, '-d', display, '-name', windowname]
        xwininfo = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                    stderr=subprocess.PIPE)
        out, err = xwininfo.communicate()
        retcode = xwininfo.wait()
        if retcode != 0:
            print "Error finding console xwindow info: " + err
            raise UninitializedWindowError

        width = None
        height = None
        mapped = False
        for line in out.splitlines():
            line = line.strip()
            if line.startswith("Map State:"):
                if line.split()[-1] != "IsViewable":
                    # Window is not mapped yet.
                    raise UnmappedWindowError
                else:
                    mapped = True
            if line.startswith("Width:"):
                width = int(line.split()[1])
            elif line.startswith("Height:"):
                height = int(line.split()[1])
            if width and height and mapped:
                return [width, height]
        else:
            # What, no width and height???
            print "No window geometry info returned by " + XWINFINFO
            raise UnmappedWindowError

    retries = 10
    sleep = 1
    uninit_count = 0
    unmap_count = 0
    width = 0
    height = 0
    while uninit_count < retries and unmap_count < retries:
        try:
            width, height = _get_window_geometry(display,
                                                 'Zone Console: ' + zonename)
            print "Discovered xterm geometry: %d x %d" % (width, height)
            break
        except UninitializedWindowError:
            if uninit_count < retries:
                print "xterm window not initialized yet. Retrying in %ds" \
                    % sleep
                uninit_count += 1
                time.sleep(sleep)
                continue
            else:
                print "xterm window is taking too long to initialize"
                break
        except UnmappedWindowError:
            if unmap_count < retries:
                print "Discovered xterm not mapped yet. Retrying in %ds" \
                    % sleep
                unmap_count += 1
                time.sleep(sleep)
                continue
            else:
                print "Discovered xterm window is taking too long to map"
                break
    else:
        print "Too many failed attempts to discover xterm window geometry"
        return

    # Generate a mode line for width and height, with a refresh of 60.0Hz
    cmd = [GTF, str(width), str(height), '60.0', '-x']
    gtf = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                           stderr=subprocess.PIPE)
    out, err = gtf.communicate()
    retcode = gtf.wait()
    if retcode != 0:
        print "Error creating new modeline for VNC display: " + err
        return

    for line in out.splitlines():
        line = line.strip()
        if line.startswith('Modeline'):
            modeline = line.split('Modeline')[1]
            print "New optimal modeline for Xvnc server: " + modeline
            mode = modeline.split()
            break

    # Create a new mode for the Xvnc server using the modeline generated by gtf
    cmd = [XRANDR, '-d', display, '--newmode']
    cmd.extend(mode)
    newmode = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    out, err = newmode.communicate()
    retcode = newmode.wait()
    if retcode != 0:
        print "Error creating new xrandr modeline for VNC display: " + err
        return

    # Add the new mode to the default display output
    modename = mode[0]
    cmd = [XRANDR, '-d', display, '--addmode', 'default', modename]
    addmode = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    out, err = addmode.communicate()
    retcode = addmode.wait()
    if retcode != 0:
        print "Error adding new xrandr modeline for VNC display: " + err
        return

    # Activate the new mode on the default display output
    cmd = [XRANDR, '-d', display, '--output', 'default', '--mode', modename]
    addmode = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    out, err = addmode.communicate()
    retcode = addmode.wait()
    if retcode != 0:
        print "Error setting new xrandr modeline for VNC display: " + err
        return

if __name__ == "__main__":
    os.putenv("LC_ALL", "C")
    smf_include.smf_main()