# 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()