#!/usr/bin/env python
# -*- coding: utf-8 -*-

# simple illustration client/server pair; client program sends a string
# to server, which echoes it back to the client (in multiple copies),
# and the latter prints to the screen
# this is the client
import socket
import sys
import struct
import re
import time
import md5

PORT = 9101
BUFSIZE = 1048576

# BNET signals
BNET_SIGNALS = {
    'BNET_EOD'           : -1,          # End of data stream, new data may follow
    'BNET_EOD_POLL'      : -2,          # End of data and poll all in one
    'BNET_STATUS'        : -3,          # Send full status
    'BNET_TERMINATE'     : -4,          # Conversation terminated, doing close()
    'BNET_POLL'          : -5,          # Poll request, I'm hanging on a read
    'BNET_HEARTBEAT'     : -6,          # Heartbeat Response requested
    'BNET_HB_RESPONSE'   : -7,          # Only response permited to HB
    'BNET_PROMPT'        : -8,          # Prompt for subcommand
    'BNET_BTIME'         : -9,          # Send UTC btime
    'BNET_BREAK'         : -10,         # Stop current command -- ctl-c
    'BNET_START_SELECT'  : -11,         # Start of a selection list
    'BNET_END_SELECT'    : -12,         # End of a select list
    'BNET_INVALID_CMD'   : -13,         # Invalid command sent
    'BNET_CMD_FAILED'    : -14,         # Command failed
    'BNET_CMD_OK'        : -15,         # Command succeeded
    'BNET_CMD_BEGIN'     : -16,         # Start command execution
    'BNET_MSGS_PENDING'  : -17,         # Messages pending
    'BNET_MAIN_PROMPT'   : -18,         # Server ready and waiting
    'BNET_SELECT_INPUT'  : -19,         # Return selection input
    'BNET_WARNING_MSG'   : -20,         # Warning message
    'BNET_ERROR_MSG'     : -21,         # Error message -- command failed
    'BNET_INFO_MSG'      : -22,         # Info message -- status line
    'BNET_RUN_CMD'       : -23,         # Run command follows
    'BNET_YESNO'         : -24,         # Request yes no response
    'BNET_START_RTREE'   : -25,         # Start restore tree mode
    'BNET_END_RTREE'     : -26,         # End restore tree mode
}

import clibbac

class ConsoleError(IOError):
    pass

class Console:
    """Class to connect with the bacula director as a console client
    """
    
    HELLO = r"Hello %s calling"
    HELLO_OK = "1000 OK"
    RXP_RESP = re.compile(
        """^auth\scram-md5(?P<compatible>c?)\s
        (?P<challenge>\S+)\s
        ssl=(?P<tls>\d)""",
        re.VERBOSE)
    
    def __init__(self, host="bacula", port=PORT):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.connect((host, port))
        self.last_state = 0
    
    def __del__(self):
        """Destructor closes the network socket"""
        self.socket.close()
        
    def _send(self, data):
        """Prepend size of data before the data string
           and send it trough the socket"""
        size = struct.pack("!i", len(data))
        self.socket.send(size + data)
    
    def _recv(self, bufsize=BUFSIZE):
        """Recieve data from socket and strip the prepended
           data size"""      
        r = self.socket.recv(4)
        size = struct.unpack("!i", r)[0]
        if size < 0:
            self.last_state = size
            r = ""
        elif size + 4 > BUFSIZE:
            raise ConsoleError("Buffer to small: %i" %size)
        else:
            r = self.socket.recv(size)
        return r
    
    def _recv_all(self):
        """Gets all lines of a request"""
        r = ""
        s = self._recv()
        while s:
            r += s
            s = self._recv()
        return r
        
    def login(self, username, password):
        pw = md5.md5(password).hexdigest()
        
        self._send(self.HELLO % username + "\n")
        response = self._recv()

        # Challenge from director
        m = self.RXP_RESP.search(response)
        challenge = clibbac.hmac_md5(m.group("challenge"), pw)
        self._send(clibbac.bin_to_base64(challenge[:16], False))
        resp = self._recv()
        if not resp.startswith(self.HELLO_OK):
            raise ConsoleError("Challenge: %s" % resp)

        # now from client site
        now = time.time()
        respond = clibbac.bin_to_base64(
            clibbac.hmac_md5("<%10.10f@bconsole>" % now, pw)[:16],
            False) 
        self._send("auth cram-md5 <%10.10f@bconsole> ssl=0\n" % now)
        resp = self._recv().rstrip("\x000")
        if resp != respond:
            print len(resp)
            print len(respond)
            raise ConsoleError("Respond: %s expected: %s" % (resp, respond))
        self._send("1000 OK auth\n")
    
        return self._recv()

    def version(self):
        """Returns the bacula director version string"""
        self._send("version")
        return self._recv_all().split("\n")[0]
    
    def dotstatus(self):
        self._send(".status dir current")
        return self._recv_all()        
    
if __name__ == "__main__":
    try:
        host = sys.argv[1]
        port = int(sys.argv[2])
    except:
        host = "bacula-devel"
        port = 9101

#    print sys.path

    console = Console(host, port)
    console.login("python", "dass-it")
        
    print console.version()
    print console.dotstatus()


#....Hello python calling
#...7auth cram-md5 <1938954300.1244114680@bacula-dir> ssl=0
#....XRx1dCJU36/UJD/l7TBZZC....
#1000 OK auth
#...4auth cram-md5 <238682207.1244114679@bconsole> ssl=0
#....klIi6CczoW89Bl/wCA04wA....
#1000 OK auth
#...31000 OK: bacula-dir Version: 3.0.1 (30 April 2009)
