#!/usr/bin/python3.12
# -*- coding: utf-8 -*-

#   Copyright (C) 2008 David Goncalves <david@lestat.st>
#   Copyright (C) 2021-2026 Jim Klimov <jimklimov+nut@gmail.com>
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.

# 2008-01-14 David Goncalves
#            PyNUT is an abstraction class to access NUT (Network UPS Tools) server.
#
# 2008-06-09 David Goncalves
#            Added 'GetRWVars' and 'SetRWVar' commands.
#
# 2009-02-19 David Goncalves
#            Changed class PyNUT to PyNUTClient
#
# 2010-07-23 David Goncalves - Version 1.2
#            Changed GetRWVars function that fails is the UPS is not
#            providing such vars.
#
# 2011-07-05 René Martín Rodríguez <rmrodri@ull.es> - Version 1.2.1
#            Added support for FSD, HELP and VER commands
#
# 2012-02-07 René Martín Rodríguez <rmrodri@ull.es> - Version 1.2.2
#            Added support for LIST CLIENTS command
#
# 2014-06-03 george2 - Version 1.3.0
#            Added custom exception class, fixed minor bug, added Python 3 support.
#
# 2021-09-27 Jim Klimov <jimklimov+nut@gmail.com> - Version 1.4.0
#            Revise strings used to be byte sequences as required by telnetlib
#            in Python 3.9, by spelling out b"STR" or str.encode('ascii');
#            the change was also tested to work with Python 2.7, 3.4, 3.5 and
#            3.7 (to the extent of accompanying test_nutclient.py at least).
#
# 2022-08-12 Jim Klimov <jimklimov+nut@gmail.com> - Version 1.5.0
#            Fix ListClients() method to actually work with current NUT protocol
#            Added DeviceLogin() method
#            Added GetUPSNames() method
#            Fixed raised PyNUTError() exceptions to carry a Python string
#            (suitable for Python2 and Python3), not byte array from protocol,
#            so exception catchers can process them naturally (see test script).
#
# 2023-01-18 Jim Klimov <jimklimov+nut@gmail.com> - Version 1.6.0
#            Added CheckUPSAvailable() method originally by Michal Hlavinka
#            from 2013-01-07 RedHat/Fedora packaging
#
# 2024-07-01 Jim Klimov <jimklimov+nut@gmail.com> - Version 1.7.0
#            Re-arranged dependency on telnetlib module (deprecated/removed
#            since Python 3.11/3.13), so we can fall back on a privately
#            stashed copy until a better solution is developed.
#
# 2025-01-31 cgar <github.com/cgarz> - Version 1.8.0
#            Removed telnetlib dependency. Switched to using socket directly.
#
# 2026-03-16 Jim Klimov <jimklimov+nut@gmail.com> - Version 1.9.0
#            Implement selectable OpenSSL and NSS client support via
#            NUT STARTTLS protocol action.

import socket

ssl_available = False
try:
    import ssl
    ssl_available = True
except:
    pass

class PyNUTError( Exception ) :
    """ Base class for custom exceptions """


class PyNUTClient :
    """ Abstraction class to access NUT (Network UPS Tools) server """

    __debug       = None   # Set class to debug mode (prints everything useful for debugging...)
    __host        = None
    __port        = None
    __login       = None
    __password    = None
    __timeout     = None
    __srv_handler = None
    __recv_leftover = b''
    # Subject to global ssl_available:
    __use_ssl     = False
    __ssl_context = None

    __version     = "1.9.0"
    __release     = "2026-03-16"

    try:
        NUT_DEFAULT_CONNECT_TIMEOUT = float(os.getenv("NUT_DEFAULT_CONNECT_TIMEOUT", "5.0"))
    except:
        NUT_DEFAULT_CONNECT_TIMEOUT = 5.0

    def __init__( self, host="127.0.0.1", port=3493, login=None, password=None, debug=False, timeout=NUT_DEFAULT_CONNECT_TIMEOUT,
                  use_ssl=False, ssl_context=None, cert_verify=None, ca_file=None, ca_path=None,
                  cert_file=None, key_file=None, key_pass=None, force_ssl=False ) :
        """ Class initialization method

host        : Host to connect (default to localhost)
port        : Port where NUT listens for connections (default to 3493)
login       : Login used to connect to NUT server (default to None for no authentication)
password    : Password used when using authentication (default to None)
debug       : Boolean, put class in debug mode (prints everything on console, default to False)
timeout     : Timeout used to wait for network response
use_ssl     : Boolean, use SSL/TLS for connection (default to False; subject to 'ssl' module availability)
ssl_context : ssl.SSLContext object to use for SSL/TLS connection (default to None; subject to 'ssl' module availability)
cert_verify : Boolean, verify server certificate (default to None, which means True if CA info is provided)
ca_file     : Path to CA certificate file (PEM format)
ca_path     : Path to directory with CA certificates (PEM format)
cert_file   : Path to client certificate file (PEM format)
key_file    : Path to client private key file (PEM format)
key_pass    : Password for client private key file (default to None)
force_ssl   : Boolean, if True, the connection must be secure (default to False)
        """
        self.__debug = debug

        if self.__debug :
            print( "[DEBUG] Class initialization..." )
            print( "[DEBUG]  -> Host  = %s (port %s, timeout %s)" % ( host, port, timeout) )
            print( "[DEBUG]  -> Login = '%s' / '%s'" % ( login, password ) )
            print( "[DEBUG]  -> SSL   = %s (force: %s, verify: %s)" % (use_ssl, force_ssl, cert_verify) )

        self.__host        = host
        self.__port        = port
        self.__login       = login
        self.__password    = password
        # TOTHINK: we could call socket.getdefaulttimeout() but it
        # might upset other scripts that use PyNUTClient, and/or
        # their imposed timeouts might confuse NUT expectations:
        self.__timeout     = NUT_DEFAULT_CONNECT_TIMEOUT if (timeout is None or type(timeout) not in [float, int] or timeout < 0.0) else timeout
        if ssl_available:
            self.__use_ssl     = use_ssl or force_ssl
            self.__force_ssl   = force_ssl
            self.__ssl_context = ssl_context
            if self.__use_ssl and self.__ssl_context is None:
                if self.__debug:
                    print( "[DEBUG] Creating SSL context" )
                # Create a default context and apply parameters
                self.__ssl_context = ssl.create_default_context(cafile=ca_file, capath=ca_path)
                if cert_verify is not None:
                    if not cert_verify:
                        # Must be set first, else we get
                        #   ValueError: Cannot set verify_mode to CERT_NONE when check_hostname is enabled.
                        self.__ssl_context.check_hostname = False
                    self.__ssl_context.verify_mode = ssl.CERT_REQUIRED if cert_verify else ssl.CERT_NONE
                if cert_file:
                    self.__ssl_context.load_cert_chain(certfile=cert_file, keyfile=key_file, password=key_pass)
                if self.__debug:
                    print( "[DEBUG] SSL context created")
        else:
            self.__use_ssl     = False
            self.__force_ssl   = False
            self.__ssl_context = None
            if (use_ssl or force_ssl or ssl_context is not None):
                if self.__debug:
                    print( "[DEBUG] SSL requested but 'ssl' module was not available for import" )
                if force_ssl:
                    raise PyNUTError( "SSL required but 'ssl' module not available" )

        self.__connect()

    # Try to disconnect cleanly when class is deleted ;)
    def __del__( self ) :
        """ Class destructor method """
        try :
            self.__srv_handler.send( b"LOGOUT\n" )
        except :
            pass

    def __read_until(self, finished_reading_data):
        data = self.__recv_leftover
        while True:
            data_end_index = data.find(finished_reading_data)
            if data_end_index == -1:
                data += self.__srv_handler.recv(50)  # nut_telnetlib.py uses 50
            else:
                break
        data_end_index += len(finished_reading_data)

        if data_end_index == len(data):
            self.__recv_leftover = b''
        else:
            self.__recv_leftover = data[data_end_index:]
            data = data[:data_end_index]
        return data

    def __connect( self ) :
        """ Connects to the defined server

If login/pass was specified, the class tries to authenticate. An error is raised
if something goes wrong.
        """
        if self.__debug :
            print( "[DEBUG] Connecting to host" )

        self.__srv_handler = socket.create_connection(
            (self.__host, self.__port),
            self.__timeout
        )

        if self.__use_ssl :
            if self.__debug :
                print( "[DEBUG] Requesting STARTTLS" )
            self.__srv_handler.send( b"STARTTLS\n" )
            result = self.__read_until( b"\n" )
            if result[:11] == b"OK STARTTLS" :
                if self.__debug :
                    print( "[DEBUG] STARTTLS accepted, wrapping socket" )
                if self.__ssl_context is None :
                    self.__ssl_context = ssl.create_default_context()
                self.__srv_handler = self.__ssl_context.wrap_socket(
                    self.__srv_handler,
                    server_hostname=self.__host
                )
            else :
                if self.__debug :
                    print( "[DEBUG] STARTTLS failed: %s" % result )
                if self.__force_ssl:
                    raise PyNUTError( "STARTTLS failed but SSL is required: %s" % result.replace( b"\n", b"" ).decode('ascii') )
                else:
                    if self.__debug:
                        print( "[DEBUG] STARTTLS failed, but not forced, continuing insecurely" )
                    self.__use_ssl = False

        if self.__login != None :
            self.__srv_handler.send( ("USERNAME %s\n" % self.__login).encode('ascii') )
            result = self.__read_until( b"\n" )
            if result[:2] != b"OK" :
                raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

        if self.__password != None :
            self.__srv_handler.send( ("PASSWORD %s\n" % self.__password).encode('ascii') )
            result = self.__read_until( b"\n" )
            if result[:2] != b"OK" :
                if result == b"ERR INVALID-ARGUMENT\n" :
                    # Quote the password (if it has whitespace etc)
                    # TODO: Escape special chard like NUT does?
                    self.__srv_handler.send( ("PASSWORD \"%s\"\n" % self.__password).encode('ascii') )
                    result = self.__read_until( b"\n" )
                    if result[:2] != b"OK" :
                        raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )
                else:
                    raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

    def GetUPSList( self ) :
        """ Returns the list of available UPS from the NUT server

The result is a dictionary containing 'key->val' pairs of 'UPSName' and 'UPS Description'

Note that fields here are byte sequences (not locale-aware strings)
which is of little concern for Python2 but is important in Python3
(e.g. when we use "str" type `ups` variables or check their "validity").
        """
        if self.__debug :
            print( "[DEBUG] GetUPSList from server" )

        self.__srv_handler.send( b"LIST UPS\n" )
        result = self.__read_until( b"\n" )
        if result != b"BEGIN LIST UPS\n":
            raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

        result = self.__read_until( b"END LIST UPS\n" )
        ups_list = {}

        for line in result.split( b"\n" ) :
            if line[:3] == b"UPS" :
                ups, desc = line[4:-1].split( b'"' )
                ups_list[ ups.replace( b" ", b"" ) ] = desc

        return( ups_list )

    def GetUPSNames( self ) :
        """ Returns the list of available UPS names from the NUT server as strings

The result is a set of str objects (comparable with ups="somename" and
useful as arguments to other methods). Helps work around Python2/Python3
string API changes.
        """
        if self.__debug :
            print( "[DEBUG] GetUPSNames from server" )

        self_ups_list = []
        for b in self.GetUPSList():
            self_ups_list.append(b.decode('ascii'))

        return self_ups_list

    def GetUPSVars( self, ups="" ) :
        """ Get all available vars from the specified UPS

The result is a dictionary containing 'key->val' pairs of all
available vars.
        """
        if self.__debug :
            print( "[DEBUG] GetUPSVars called..." )

        self.__srv_handler.send( ("LIST VAR %s\n" % ups).encode('ascii') )
        result = self.__read_until( b"\n" )
        if result != ("BEGIN LIST VAR %s\n" % ups).encode('ascii') :
            raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

        ups_vars   = {}
        result     = self.__read_until( ("END LIST VAR %s\n" % ups).encode('ascii') )
        offset     = len( ("VAR %s " % ups ).encode('ascii') )
        end_offset = 0 - ( len( ("END LIST VAR %s\n" % ups).encode('ascii') ) + 1 )

        for current in result[:end_offset].split( b"\n" ) :
            var  = current[ offset: ].split( b'"' )[0].replace( b" ", b"" )
            data = current[ offset: ].split( b'"' )[1]
            ups_vars[ var ] = data

        return( ups_vars )

    def CheckUPSAvailable( self, ups="" ) :
        """ Check whether UPS is reachable

Just tries to contact UPS with a safe command.
The result is True (reachable) or False (unreachable)
        """
        if self.__debug :
            print( "[DEBUG] CheckUPSAvailable called..." )

        self.__srv_handler.send( ("LIST CMD %s\n" % ups).encode('ascii') )
        result = self.__read_until( b"\n" )
        if result != ("BEGIN LIST CMD %s\n" % ups).encode('ascii') :
            return False

        self.__read_until( ("END LIST CMD %s\n" % ups).encode('ascii') )
        return True

    def GetUPSCommands( self, ups="" ) :
        """ Get all available commands for the specified UPS

The result is a dict object with command name as key and a description
of the command as value
        """
        if self.__debug :
            print( "[DEBUG] GetUPSCommands called..." )

        self.__srv_handler.send( ("LIST CMD %s\n" % ups).encode('ascii') )
        result = self.__read_until( b"\n" )
        if result != ("BEGIN LIST CMD %s\n" % ups).encode('ascii') :
            raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

        ups_cmds   = {}
        result     = self.__read_until( ("END LIST CMD %s\n" % ups).encode('ascii') )
        offset     = len( ("CMD %s " % ups).encode('ascii') )
        end_offset = 0 - ( len( ("END LIST CMD %s\n" % ups).encode('ascii') ) + 1 )

        for current in result[:end_offset].split( b"\n" ) :
            var  = current[ offset: ].split( b'"' )[0].replace( b" ", b"" )

            # For each var we try to get the available description
            try :
                self.__srv_handler.send( ("GET CMDDESC %s %s\n" % ( ups, var )).encode('ascii') )
                temp = self.__read_until( b"\n" )
                if temp[:7] != b"CMDDESC" :
                    raise PyNUTError
                else :
                    off  = len( ("CMDDESC %s %s " % ( ups, var )).encode('ascii') )
                    desc = temp[off:-1].split(b'"')[1]
            except :
                desc = var

            ups_cmds[ var ] = desc

        return( ups_cmds )

    def GetRWVars( self,  ups="" ) :
        """ Get a list of all writable vars from the selected UPS

The result is presented as a dictionary containing 'key->val' pairs
        """
        if self.__debug :
            print( "[DEBUG] GetUPSVars from '%s'..." % ups )

        self.__srv_handler.send( ("LIST RW %s\n" % ups).encode('ascii') )
        result = self.__read_until( b"\n" )
        if ( result != ("BEGIN LIST RW %s\n" % ups).encode('ascii') ) :
            raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

        result     = self.__read_until( ("END LIST RW %s\n" % ups).encode('ascii') )
        offset     = len( ("VAR %s" % ups).encode('ascii') )
        end_offset = 0 - ( len( ("END LIST RW %s\n" % ups).encode('ascii') ) + 1 )
        rw_vars    = {}

        try :
            for current in result[:end_offset].split( b"\n" ) :
                var  = current[ offset: ].split( b'"' )[0].replace( b" ", b"" )
                data = current[ offset: ].split( b'"' )[1]
                rw_vars[ var ] = data

        except :
            pass

        return( rw_vars )

    def SetRWVar( self, ups="", var="", value="" ):
        """ Set a variable to the specified value on selected UPS

The variable must be a writable value (cf GetRWVars) and you must have the proper
rights to set it (maybe login/password).
        """

        self.__srv_handler.send( ("SET VAR %s %s %s\n" % ( ups, var, value )).encode('ascii') )
        result = self.__read_until( b"\n" )
        if ( result == b"OK\n" ) :
            return( "OK" )
        else :
            raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

    def RunUPSCommand( self, ups="", command="" ) :
        """ Send a command to the specified UPS

Returns OK on success or raises an error
        """

        if self.__debug :
            print( "[DEBUG] RunUPSCommand called..." )

        self.__srv_handler.send( ("INSTCMD %s %s\n" % ( ups, command )).encode('ascii') )
        result = self.__read_until( b"\n" )
        if ( result == b"OK\n" ) :
            return( "OK" )
        else :
            raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

    def DeviceLogin( self, ups="") :
        """ Establish a login session with a device (like upsmon does)

Returns OK on success or raises an error
USERNAME and PASSWORD must have been specified earlier in the session (once)
and upsd.conf should permit that user with one of `upsmon` role types.

Note there is no "device LOGOUT" in the protocol, just one for general end
of connection.
        """

        if self.__debug :
            print( "[DEBUG] DeviceLogin called..." )

        if ups is None or (ups not in self.GetUPSNames()):
            if self.__debug :
                print( "[DEBUG] DeviceLogin: %s is not a valid UPS" % ups )
            raise PyNUTError( "ERR UNKNOWN-UPS" )

        self.__srv_handler.send( ("LOGIN %s\n" % ups).encode('ascii') )
        result = self.__read_until( b"\n" )
        if ( result.startswith( ("User %s@" % self.__login).encode('ascii')) and result.endswith (("[%s]\n" % ups).encode('ascii')) ):
            # User dummy-user@127.0.0.1 logged into UPS [dummy]
            # Read next line then
            result = self.__read_until( b"\n" )
        if ( result == b"OK\n" ) :
            return( "OK" )
        else :
            raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

    def FSD( self, ups="") :
        """ Send FSD command

Returns OK on success or raises an error

NOTE: API changed since NUT 2.8.0 to replace MASTER with PRIMARY
(and backwards-compatible alias handling)
        """

        if self.__debug :
            print( "[DEBUG] PRIMARY called..." )

        self.__srv_handler.send( ("PRIMARY %s\n" % ups).encode('ascii') )
        result = self.__read_until( b"\n" )
        if ( result != b"OK PRIMARY-GRANTED\n" ) :
            if self.__debug :
                print( "[DEBUG] Retrying: MASTER called..." )
            self.__srv_handler.send( ("MASTER %s\n" % ups).encode('ascii') )
            result = self.__read_until( b"\n" )
            if ( result != b"OK MASTER-GRANTED\n" ) :
                if self.__debug :
                    print( "[DEBUG] Primary level functions are not available" )
                raise PyNUTError( "ERR ACCESS-DENIED" )

        if self.__debug :
            print( "[DEBUG] FSD called..." )
        self.__srv_handler.send( ("FSD %s\n" % ups).encode('ascii') )
        result = self.__read_until( b"\n" )
        if ( result == b"OK FSD-SET\n" ) :
            return( "OK" )
        else :
            raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

    def help(self) :
        """ Send HELP command
        """

        if self.__debug :
            print( "[DEBUG] HELP called..." )

        self.__srv_handler.send( b"HELP\n" )
        return self.__read_until( b"\n" )

    def ver(self) :
        """ Send VER command
        """

        if self.__debug :
            print( "[DEBUG] VER called..." )

        self.__srv_handler.send( b"VER\n" )
        return self.__read_until( b"\n" )

    def ListClients( self, ups = None ) :
        """ Returns the list of connected clients from the NUT server

The result is a dictionary containing 'key->val' pairs of 'UPSName' and a list of clients
        """
        if self.__debug :
            print( "[DEBUG] ListClients from server: %s" % ups )

        # If (!ups) we use this list below to recurse:
        self_ups_list = self.GetUPSNames()
        if ups and (ups not in self_ups_list):
            if self.__debug :
                print( "[DEBUG] ListClients: %s is not a valid UPS" % ups )
            raise PyNUTError( "ERR UNKNOWN-UPS" )

        if ups:
            self.__srv_handler.send( ("LIST CLIENT %s\n" % ups).encode('ascii') )
        else:
            # NOTE: Currently NUT does not support just listing all clients
            # (not providing an "ups" argument) => NUT_ERR_INVALID_ARGUMENT
            self.__srv_handler.send( b"LIST CLIENT\n" )
        result = self.__read_until( b"\n" )
        if ( (ups and result != ("BEGIN LIST CLIENT %s\n" % ups).encode('ascii')) or (ups is None and result != b"BEGIN LIST CLIENT\n") ):
            if ups is None and (result == b"ERR INVALID-ARGUMENT\n") :
                # For ups==None, list all upses, list their clients
                if self.__debug :
                    print( "[DEBUG] Recurse ListClients() because it did not specify one UPS to query" )
                ups_list = {}
                for ups in self_ups_list :
                    # Update "ups_list" dict with contents of recursive call return
                    ups_list.update(self.ListClients(ups))
                return( ups_list )

            # had a seemingly valid arg, but no success:
            if self.__debug :
                print( "[DEBUG] ListClients from server got unexpected result: %s" % result )

            raise PyNUTError( result.replace( b"\n", b"" ).decode('ascii') )

        if ups:
            result = self.__read_until( ("END LIST CLIENT %s\n" % ups).encode('ascii') )
        else:
            # Should not get here with current NUT:
            result = self.__read_until( b"END LIST CLIENT\n" )
        ups_list = {}

        for line in result.split( b"\n" ):
            ###print( "[DEBUG] ListClients line: '%s'" % line )
            if line[:6] == b"CLIENT" :
                ups, host = line[7:].split(b' ')
                ups.replace(b' ', b'')
                host.replace(b' ', b'')
                if not ups in ups_list:
                    ups_list[ups] = []
                ups_list[ups].append(host)

        return( ups_list )
