#!/usr/bin/python2

import socket
from ctypes import *
import struct

PACKET_HELLO   = 0
PACKET_AUTH    = 1
PACKET_REBOOT  = 2
PACKET_SUCCESS = 254
PACKET_FAIL    = 255


class packet(BigEndianStructure):
    _pack_ = 1
    _fields_ = [("magic", c_int64),  # The magic word
                ("token", c_uint32), # Unique auth token, per-session
                ("type", c_int16),   # Type of packet
                ("size", c_uint32)]  # Size of the packet body

    def get_type(self):
        if self.type == PACKET_HELLO:
            return "PACKET_HELLO"
        elif self.type == PACKET_AUTH:
            return "PACKET_AUTH"
        elif self.type == PACKET_REBOOT:
            return "PACKET_REBOOT"
        elif self.type == PACKET_SUCCESS:
            return "PACKET_SUCCESS"
        elif self.type == PACKET_FAIL:
            return "PACKET_FAIL"
        else:
            return "PACKET_UNKNOWN"


class pkt_auth(BigEndianStructure):
    _pack_ = 1
    _fields_ = [("login_size", c_uint32),
                ("password_size", c_uint32)] # append login+password


class pkt_reboot(BigEndianStructure):
    _pack_ = 1
    _fields_ = [] # append message


class JurassicProtocol(object):
    def __init__(self, host, port, magic):
        self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.s.connect((host, port))
        self.token = 0
        self.magic = magic
        self.p = None    # Last packet
        self.body = None # Last packet body
        self.get_token()

    def __del__(self):
        if self.s:
            self.s.close()

    def recv_len(self, size):
        res = ''

        while not size == 0:
            tmp = self.s.recv(size)
            if tmp == "":
                raise Exception
            res += tmp
            size -= len(tmp)

        return res

    def read_packet(self):
        res = self.s.recv(sizeof(packet))
        self.p = packet.from_buffer(bytearray(res))

        if self.p.size:
            self.body = bytearray(self.s.recv(self.p.size))
        else:
            self.body = ""

    def send_packet(self, type, body):
        self.p = packet(self.magic, self.token, type, len(body))
        self.s.sendall(self.p)
        self.s.sendall(body)

    def get_token(self):
        self.read_packet()
        self.token = self.p.token

    def authenticate(self, login, password):
        body = bytearray(pkt_auth(len(login), len(password))) + login + password
        self.send_packet(PACKET_AUTH, body)
        self.read_packet()

        if len(self.body) == 4:
            return (struct.unpack('>I', self.body)[0], "")
        else:
            return (struct.unpack('>I', self.body[:4])[0], self.body[4:])

    def reboot(self, message):
        body = bytearray(pkt_reboot()) + message
        self.send_packet(PACKET_REBOOT, body)
        self.read_packet()
        return self.body

    def dump_packet(self):
        print "Magic: %#x" % self.p.magic
        print "Token: %#x" % self.p.token
        print "Type:", self.p.get_type()
        print "Body size:", self.p.size
        print "Body:", repr(self.body)
