#!/usr/bin/python2.7
# -*- coding: utf-8 -*-

__author__ = 'dal'

from pexpect import TIMEOUT, EOF
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *
import argparse
import commands
import newrole_gui
import os
import pexpect
import sys

_ = newrole_gui.getTransFunc()


class MainDialog(QDialog):

    def __init__(self, parent = None):
        super(MainDialog, self).__init__(parent)

        self._CurrentSelinuxContext = newrole_gui.getCurrentContext()
        self.setWindowTitle( _( 'Run programm'.encode( 'utf-8' ) ) )
        self.setMinimumSize( 400, self.minimumSize().height() )
        #self.setMaximumSize(450, self.maximumSize().height())

        self._layout = QVBoxLayout()
        self.setLayout( self._layout )

        self._devicesLabel     = QLabel( _('Command:') )
        self._layout.addWidget( self._devicesLabel )
        self._sCommandComboBox = QComboBox()
        self._sCommandComboBox.setEditable( True )
        self._sCommandComboBox.editTextChanged.connect( self.checkCommand )
        self._layout.addWidget( self._sCommandComboBox )

        self._contextLabel  = QLabel( _('Context:') )
        self._layout.addWidget( self._contextLabel )

        self._contextLayout = QHBoxLayout()
        self._layout.addLayout( self._contextLayout )

        self._sRoleComboBox = QComboBox()
        self._sRoleComboBox.setEditable( True )
        self._sRoleComboBox.editTextChanged.connect( self.checkRole )
        self._contextLayout.addWidget( self._sRoleComboBox )

        self._sTypeComboBox = QComboBox()
        self._sTypeComboBox.setEditable( True )
        self._sTypeComboBox.editTextChanged.connect( self.checkType )
        self._sTypeComboBox.setSizePolicy( QSizePolicy.Minimum, QSizePolicy.Preferred )
        self._contextLayout.addWidget( self._sTypeComboBox )

        self._sLevelLineEdit = QLineEdit( 's0' )
        self._sLevelLineEdit.setMinimumWidth( 50 )
        self._sLevelLineEdit.resize( 20, self._sLevelLineEdit.height() )
        self._sLevelLineEdit.textChanged.connect( self.checkLevel )
        self._sLevelLineEdit.setSizePolicy( QSizePolicy.Minimum, QSizePolicy.Preferred )
        self._contextLayout.addWidget( self._sLevelLineEdit )

        self._buttonBox    = QDialogButtonBox()
        self._cancelButton = QPushButton( _('Cancel') )
        self._buttonBox.addButton( self._cancelButton, QDialogButtonBox.RejectRole )
        self._runButton    = QPushButton( _('Run') )
        self._buttonBox.addButton( self._runButton, QDialogButtonBox.AcceptRole )
        self._layout.addWidget(self._buttonBox )

        self._buttonBox.accepted.connect( self.accept )
        self._buttonBox.rejected.connect( self.reject )

        self.resize( 450, self.minimumHeight() )

    def checkLevel(self):
        # It's very difficult to make regex, so just try to parse
        level = self._sLevelLineEdit.text()

        if level == '':
            self.resetWarning( self._sLevelLineEdit )
            return

        if not newrole_gui.validateLevel( level ):
            self.showContextWarning( self._sLevelLineEdit, _('Bad format') )
        else:

            self.resetWarning( self._sLevelLineEdit )
        return

    def checkRole(self):

        role  = self._sRoleComboBox.currentText()
        types = self.seinfoProvider.getTypes( role )

        if types:
            self.FillCombobox( self._sTypeComboBox, types )

        index = self._sRoleComboBox.findText( role )

        if index<0:
            self.showContextWarning( self._sRoleComboBox, _('Unknown role') )
        else:
            self.resetWarning( self._sRoleComboBox )

        deftype = self.seinfoProvider.getDefaultType( role )
        index   = self._sTypeComboBox.findText( deftype )
        self._sTypeComboBox.setCurrentIndex( index )
        self.checkType()

        return

    def checkType(self):

        role  = self._sTypeComboBox.currentText()
        index = self._sTypeComboBox.findText( role )

        if index<0:
            self.showContextWarning( self._sTypeComboBox, _('Unknown type') )
        else:
            self.resetWarning( self._sTypeComboBox )

        return

    def checkCommand(self):
        return

    def resetWarning(self, elem):

        Palette = QApplication.palette()
        elem.setPalette( Palette )

    def showContextWarning(self, elem, text):

        warningPalette = QPalette()
        warningPalette.setColor( QPalette.Base, Qt.yellow )
        elem.setPalette( warningPalette )
        QToolTip.showText( elem.mapToGlobal( QPoint( 0, 0 ) ), text, elem )

    def setCommandInfo(self, commands, command):

        self._sCommandComboBox.addItems( commands )
        self._sCommandComboBox.setEditText( command )

    def FillCombobox(self, combobox, list):
        combobox.my_completer = QCompleter( list, self )
        combobox.my_completer.setCaseSensitivity( Qt.CaseInsensitive )
        combobox.clear()
        combobox.addItems( list )
        combobox.setCompleter( combobox.my_completer )

class Ask(QDialog):

    def __init__(self, title='', pin=False, parent = None):
        super( Ask, self ).__init__( parent )

        if title:
            self.setWindowTitle( title )
        else:
            if pin:
                self.setWindowTitle( _('PIN required') )
            else:
                self.setWindowTitle( _('Password required') )

        self.setMinimumSize( 400, self.minimumSize().height() )

        self._VLayout = QVBoxLayout()
        self.setLayout( self._VLayout )


        self._HLayout = QHBoxLayout()
        self._Fields  = QWidget()
        self._Fields.setLayout( self._HLayout )

        if pin:
            self._PinLabel     = QLabel( _('PIN:') )
            self._HLayout.addWidget( self._PinLabel )
            self._sPinLineEdit = QLineEdit()
            self._sPinLineEdit.setEchoMode( QLineEdit.Password )
            self._sPinLineEdit.setMinimumWidth( 50 )
            self._HLayout.addWidget( self._sPinLineEdit )
        else:
            self._PasswordLabel     = QLabel( _('Password:') )
            self._HLayout.addWidget( self._PasswordLabel )
            self._sPasswordLineEdit = QLineEdit()
            self._sPasswordLineEdit.setEchoMode( QLineEdit.Password )
            self._sPasswordLineEdit.setMinimumWidth( 50 )
            self._HLayout.addWidget( self._sPasswordLineEdit )

        self._VLayout.addWidget( self._Fields )

        self._buttonBox    = QDialogButtonBox()
        self._cancelButton = QPushButton( _('Cancel') )
        self._buttonBox.addButton( self._cancelButton, QDialogButtonBox.RejectRole )
        self._addButton    = QPushButton( _('OK') )
        self._buttonBox.addButton( self._addButton, QDialogButtonBox.AcceptRole )
        self._VLayout.addWidget( self._buttonBox )

        self._buttonBox.accepted.connect( self.accept )
        self._buttonBox.rejected.connect( self.reject )

        self.setFixedSize( 450, self.minimumHeight() )

class resultEnum:
    OK = 0
    BAD_CONTEXT = 1
    PRINTUSAGE = 2
    CANCEL = 3
    KEYBOARDINTERRUPT = 4
    UNEXPECTEDEOF = 5

class Runner():
    def __init__(self):
        self.app = None
        self._dialogCreated = False
        self._CurrentSelinuxContext= newrole_gui.getCurrentContext()

    def getExecs(self):
        execs=[]
        # The following search is incompatible with selinux enforcing - disabled for now
        # paths = os.environ['PATH'].split(':')
        # for p in paths:
        #     ret,out = commands.getstatusoutput(r'find -L %s -executable -type f -printf "%%f\n"' % p)
        #         if ret == 0:
        #             execs+=out.split()
        execs.append('')
        execs.sort()
        return execs

    #helpers
    def insertIfNewName(self, struct, name, template):
        ins         = template
        ins['name'] = name
        if len(struct) > 0:
            for s in struct:
                if s.has_key( 'name' ):
                    if s['name'] == name:
                        return s
                    else:
                        struct.append( ins )
                        return struct[-1]
        else:
            struct.append( ins )
            return struct[-1]

    def insertIfUnique(self, list, elem):
        if elem not in list:
            list.append( elem )

    def createMainDialog(self,context = ['','',''], ItemVisibility = [True,True,True,True], command = '', set_contexts = []):

        self._dialog              = MainDialog()
        self._dialog.set_contexts = set_contexts

        if ItemVisibility[0]: self._dialog._sRoleComboBox.setEnabled( False )
        if ItemVisibility[1]: self._dialog._sTypeComboBox.setEnabled( False )
        if ItemVisibility[2]: self._dialog._sLevelLineEdit.setEnabled( False )
        if ItemVisibility[3]: self._dialog._sCommandComboBox.setEnabled( False )

        self._dialog._sRolesFull = []
        if set_contexts != None and len(set_contexts) > 0:
            roles = []

            for ctx in set_contexts:
                if ctx[0] == newrole_gui.getCurrentContext()[0]:
                    self.insertIfUnique( roles, ctx[1] )
                    role = self.insertIfNewName( self._dialog._sRolesFull, ctx[1], { 'dominate': 'not using',  'types': [] } )
                    self.insertIfUnique( role['types'], ctx[2] )

            self.extendForEmptyRole( self._dialog._sRolesFull )
            roles.insert( 0, '' )
            self._dialog.seinfoProvider = newrole_gui.ManualSeinfoProvider( ( roles, self._dialog._sRolesFull ) )

        else:
            self._dialog.seinfoProvider = newrole_gui.RealSeinfoProvider( None )

        self._sRoles = self._dialog.seinfoProvider.getRoles( newrole_gui.getCurrentContext()[0] )
        self._sTypes = self._dialog.seinfoProvider.getTypes( context[0] )

        self._dialog.FillCombobox( self._dialog._sRoleComboBox, self._sRoles )
        self._dialog.FillCombobox( self._dialog._sTypeComboBox, self._sTypes )

        self._sTypes = self._dialog.seinfoProvider.getTypes( context[0] )
        self._dialog._sRoleComboBox.setEditText( context[0] )
        self._dialog._sTypeComboBox.setEditText( context[1] )
        self._dialog._sLevelLineEdit.setText( context[2] )
        self.execs   = self.getExecs()
        self._dialog.setCommandInfo( self.execs, command )
        self._dialog.show()

    def showMainDialog(self,context = ['','',''], ItemVisibility = [True,True,True,True], command = '', set_contexts =[]):

        if  not self._dialogCreated:
            self.createMainDialog( context, ItemVisibility, command, set_contexts )

        self._dialogCreated = True

        return self._dialog.exec_()

    def extendForEmptyRole(self, _sRolesFull):

        types_for_current_role = None

        for r in _sRolesFull:
            if r['name'] == self._CurrentSelinuxContext[1]:
                types_for_current_role = r['types']

            r['types'].append('')

        if types_for_current_role:
            _sRolesFull.append( { 'dominate': [''], 'name': '', 'types':types_for_current_role } )

    def ask(self, title, pin=False):

        if pin:
            askp = Ask( title, True )
            askp.show()
            r = askp.exec_()
            return askp._sPinLineEdit.text(), r

        else:
            askp = Ask( title, False )
            askp.show()
            r = askp.exec_()
            return askp._sPasswordLineEdit.text(), r

    def childTimeOutWait(self, child):

        if child.isalive():
            child.wait()

        if child.exitstatus != 0:
            return resultEnum.UNEXPECTEDEOF, child.before

        else:
            return resultEnum.OK, 'OK'

    # Here is pexpect logic for newrole
    def execute(self, cmd, title=''):

        incorrect          = False
        token_not_found    = False
        running            = False

        cmd                = cmd.replace( '{LANG}', os.environ['LANG'] )
        env                = os.environ.copy()
        env['LC_MESSAGES'] = 'en_US.UTF-8'
        child = pexpect.spawn( cmd, env=env )

        # init password
        i = child.expect ( [ TIMEOUT,
                             'Password: ',
                             'is not a valid context',
                             "Couldn't get default type.",
                             'USAGE: newrole \[ -r role \] \[ -t type \] \[ -l level \] \[ -p \] \[ -V \] \[ -- args \]',
                              EOF
                            ], timeout = 1)
        if i == 0:
            pass

        if i == 1:
            pas, r = self.ask( title )

            if r == 0: #cancel
                return resultEnum.CANCEL, 'OK'

            child.sendline( '%s\r' % pas )
            pas = ''

        if i == 2 or i == 3:
            return resultEnum.BAD_CONTEXT, _('SELinux context is wrong.')
        if i == 4:
            return resultEnum.PRINTUSAGE, _('Print usage')
        if i == 5:
            return resultEnum.UNEXPECTEDEOF, child.before

        # init pin
        i = child.expect( [ TIMEOUT, 'Smart card PIN:', 'No suitable token available', EOF ], timeout=1 )

        if i == 0:
            pass

        if i == 1:
            pin, r = self.ask( title, True )

            if r == 0: #cancel
                return resultEnum.CANCEL, 'OK'

            child.sendline ( '%s\r' % pin )
            pin = ''

        if i == 2:
            token_not_found = True

        # check
        i = child.expect( [ 'incorrect password', EOF, TIMEOUT ] )

        if i == 0:
            incorrect = True
        if i == 1 or i == 2:
            return self.childTimeOutWait(child)

        child.expect(EOF)
        child.wait()
        child.close()

        if incorrect:
            self.execute(cmd, _('Bad password or PIN') )

        # if incorrect_pin:
        #     self.execute(cmd, _('Bad pin') )

        if token_not_found:
            self.execute(cmd, _('Token not found') )

        return resultEnum.OK, 'OK'


if __name__ == '__main__':

    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument('-r','--role',  dest='role',  action='store', help='selinux role')
    arg_parser.add_argument('-t','--type',  dest='type',  action='store', help='selinux type')
    arg_parser.add_argument('-l','--level', dest='level', action='store', help='selinux level')
    arg_parser.add_argument('-p','--preserve-environment', dest='preserve_environment', action='store_true', default=False,
                            help='preserve environment')
    arg_parser.add_argument('-c','--command', dest='command', action='store', help='command to run')
    arg_parser.add_argument('-sc','--set-contexts', dest='set_contexts', action='store', help='set contexts in gui')
    arg_parser.add_argument('-g','--gui', action='store_true', default=False, help='run GUI')

    cmdLineOptions  = arg_parser.parse_args(sys.argv[1:])

    set_contexts = []

    if cmdLineOptions.set_contexts:
        try:
            ctxs = cmdLineOptions.set_contexts.split()
            for ctx in ctxs:
                set_contexts.append( ctx.split(':') )
                if len( set_contexts[-1] ) < 4:
                    raise Exception( 'Bad syntax' )
        except :
             QMessageBox.critical( None, _('Error!'), _('Bad --set-context format.') )
             sys.exit( -1 )

    runner = Runner()

    while True:
        if cmdLineOptions.gui:
            #runner._dialog._sRoleComboBox.setEditText(cmdLineOptions.role)

            visibility = []
            visibility.append( cmdLineOptions.role    )
            visibility.append( cmdLineOptions.type    )
            visibility.append( cmdLineOptions.level   )
            visibility.append( cmdLineOptions.command )
            visibility = [0,0,0,0]
            context = []
            context.append( cmdLineOptions.role )
            context.append( cmdLineOptions.type )
            context.append( cmdLineOptions.level )

            if runner.showMainDialog( context, visibility, cmdLineOptions.command, set_contexts ):
                cmdLineOptions.role    = runner._dialog._sRoleComboBox.currentText()
                cmdLineOptions.type    = runner._dialog._sTypeComboBox.currentText()
                cmdLineOptions.level   = runner._dialog._sLevelLineEdit.text()
                cmdLineOptions.command = runner._dialog._sCommandComboBox.currentText()
            else:
                sys.exit(0)
        else:
            if not cmdLineOptions.command:
                print( 'you should specify a command' )
                break

        cmd = 'newrole'
        if cmdLineOptions.role:
            cmd += ' -r %s' % cmdLineOptions.role
        if cmdLineOptions.type:
            cmd += ' -t %s' % cmdLineOptions.type
        if cmdLineOptions.level:
            cmd += ' -l %s' % cmdLineOptions.level
        if cmdLineOptions.preserve_environment:
            cmd += ' -p '
        if cmdLineOptions.command:
            cmd += " -- -c 'LANG=\"{LANG}\" %s'" % cmdLineOptions.command
        else:
            QMessageBox.critical( None, _('Error!'), _("Command can't be empty.") )
            continue

        print( 'Executing: %s' % cmd )
        try:
            ret,desc = runner.execute( cmd )
        except KeyboardInterrupt:
            sys.exit( resultEnum.KEYBOARDINTERRUPT )

        if ret == resultEnum.OK:
            sys.exit(0)

        if ret == resultEnum.UNEXPECTEDEOF:
           QMessageBox.critical( None, _('Error!'), desc )

           if not cmdLineOptions.gui:
                sys.exit( ret )

        if ret == resultEnum.BAD_CONTEXT:
            QMessageBox.critical( None, _('Error!'), desc )

            if not cmdLineOptions.gui:
                sys.exit( resultEnum.BAD_CONTEXT )

        if ret == resultEnum.PRINTUSAGE:
            QMessageBox.critical( None, _('Error!'), _('Need more parameters to be present.') )

            if not cmdLineOptions.gui:
                sys.exit(ret)

        if ret == resultEnum.CANCEL:
            if not cmdLineOptions.gui:
                sys.exit( 0 )
    sys.exit( 0 )
