supwisdom/agentlink/__init__.py

# -*- coding: utf-8
"""
agentlink 模块，提供与 swagent 通讯协议的封装
"""
import codecs
import socket
import struct
import random
import time
import os
import traceback

__all__ = ['VERSION', 'SWAgentLink']

VERSION = "4.0.0"


class SWAgentLink(object):
    """
    agentlink 协议封装
    """
    MSG_LEN_FMT = '!H'
    MSG_HEADER_FMT = '!IIIBB'
    MSG_HEADER_DESC = ('major_version', 'minor_version', 'unique_no', 'msg_type', 'retcode')
    # 协议版本号定义
    MAJOR_VERSION = 2
    MINOR_VERSION = 1

    def __init__(self):
        self.sock = None
        self.ip = None
        self.port = None
        self.timeout = 10.0
        self._is_login = False
        self.unique_no = None
        self.connection_id = None
        random.seed(time.time())

    def connect(self, ip, port):
        self.close()
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(self.timeout)
            sock.connect((ip, port))
            self.sock = sock
            self.ip = ip
            self.port = port
            return True
        except socket.error:
            return False

    def login(self):
        msg_type = 0x30
        process_id = struct.pack('!I', os.getpid())
        message_header = self._new_message(msg_type)
        login_data = self._pack_request(message_header, process_id)
        data = self._send_and_recv(login_data, 5.0)
        if not data:
            return False
        login_resp = self._unpack_request(data)
        if not self._check_response(message_header, login_resp):
            print u"Login response error"
            return False
        if login_resp['retcode'] != 0:
            print u"Login error retcode={0}".format(login_resp['retcode'])
            return False
        self.connection_id = login_resp['data']
        self._is_login = True
        return True

    def _send_and_recv(self, request, timeout):
        header_len = struct.calcsize(self.MSG_LEN_FMT)
        try:
            self.sock.settimeout(timeout)
            self.sock.send(request)
            # print u"Send[{0}]".format(codecs.encode(request, 'hex'))
            header = self.sock.recv(header_len)
            if not header:
                return None
            body_len = struct.unpack(self.MSG_LEN_FMT, header)[0]
            body = self.sock.recv(body_len)
            # print u"Recv[{0}]".format(codecs.encode(body, 'hex'))
            return body
        except socket.error as e:
            print u"socket send and recv error, error: <{0}>".format(e)
            traceback.print_exc()
            self.sock.close()
            self.sock = None
            self._is_login = False
            return None

    def get_data(self):
        msg_type = 0x32
        message_header = self._new_message(msg_type)
        req_data = self._pack_request(message_header, self.connection_id)
        data = self._send_and_recv(req_data, None)
        if not data:
            return False
        getdata_resp = self._unpack_request(data)
        if not self._check_response(message_header, getdata_resp):
            print u"get data error"
            return None
        if getdata_resp['retcode'] != 0:
            print u"GetData error retcode={0}".format(getdata_resp['retcode'])
            return None
        data = getdata_resp['data']
        # print codecs.encode(data, 'hex')
        data_head_len = struct.calcsize('!I')
        data_len = struct.unpack('!I', data[:data_head_len])[0]
        if data_len != len(data) - data_head_len:
            return None
        return data[data_head_len: data_head_len + data_len]

    def answer_data(self, response):
        msg_type = 0x34
        message_header = self._new_message(msg_type)
        req_data = self._pack_request(message_header, self.connection_id + response)
        data = self._send_and_recv(req_data, None)
        answer_resp = self._unpack_request(data)
        if not self._check_response(message_header, answer_resp):
            print u"answer data error"
            return False
        if answer_resp['retcode'] != 0:
            print u"Answer error retcode={0}".format(answer_resp['retcode'])
            return False
        return True

    def answer_invalidate(self):
        msg_type = 0x36
        message_header = self._new_message(msg_type)
        req_data = self._pack_request(message_header, self.connection_id)
        data = self._send_and_recv(req_data, None)
        answer_resp = self._unpack_request(data)
        if not self._check_response(message_header, answer_resp):
            print u"answer data error"
            return False
        if answer_resp['retcode'] != 0:
            print u"Answer error retcode={0}".format(answer_resp['retcode'])
            return False
        return True

    def _new_message(self, msg_type):
        unique_no = random.randint(0x10000000, 0xFFFFFFFF)
        return dict(major_version=self.MAJOR_VERSION, minor_version=self.MINOR_VERSION,
                    unique_no=unique_no, msg_type=msg_type, retcode=0)

    def _pack_request(self, message_header, data):
        header = struct.pack(self.MSG_HEADER_FMT, message_header['major_version'],
                             message_header['minor_version'], message_header['unique_no'],
                             message_header['msg_type'], message_header['retcode'])
        r = []
        if not data:
            r.append(struct.pack(self.MSG_LEN_FMT, len(header)))
            r.append(data)
        else:
            r.append(struct.pack(self.MSG_LEN_FMT, (len(header) + len(data))))
            r.append(header)
            r.append(data)
        return ''.join(r)

    def close(self):
        if self.sock:
            self.sock.close()
            self.sock = None
        self._is_login = False

    def is_ok(self):
        if not self.sock:
            return False
        if not self._is_login:
            return False
        return True

    def _unpack_request(self, data):
        if not data:
            return None
        header_len = struct.calcsize(self.MSG_HEADER_FMT)
        if header_len > len(data):
            return None
        resp = struct.unpack(self.MSG_HEADER_FMT, data[:header_len])
        resp_dict = dict(zip(self.MSG_HEADER_DESC, resp))
        resp_dict['data'] = data[header_len:]
        return resp_dict

    def _check_response(self, message_header, response):
        if not response:
            return False
        if (message_header['major_version'] != response['major_version']
            or message_header['minor_version'] != response['minor_version']
            or message_header['unique_no'] != response['unique_no']
            or message_header['msg_type'] != response['msg_type']):
            return False
        return True