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

from __future__ import print_function

import hashlib
import socket
import struct
import sys

PROTO_HEADER_FMT = '<HHHI'
PROTO_HEADER_SIZE = struct.calcsize(PROTO_HEADER_FMT)
VERIFY_HEAD_FMT = '<H32s'
VERIFY_HEAD_SIZE = struct.calcsize(VERIFY_HEAD_FMT)
REQ_BODY_FMT = '<LLLLL'
REQ_BODY_SIZE = struct.calcsize(REQ_BODY_FMT)
ACK_SUCCESS_BODY_FMT = '<LL'
ACK_SUCCESS_BODY_SIZE = struct.calcsize(ACK_SUCCESS_BODY_FMT)
SEQ_NUM = 0
TIMEOUT_SEC = 5
HOST = '203.73.22.167'
PORT = 12330
CMD_ID = 1001
CHANNEL_ID = 101
CHANNEL_KEY = '12345678'


def recv_all(sock, size):
    chunks = []
    received = 0
    while received < size:
        chunk = sock.recv(size - received)
        if not chunk:
            raise RuntimeError('socket closed')
        chunks.append(chunk)
        received += len(chunk)
    return b''.join(chunks)


def build_request(cmd_id, account_id):
    dest_id = account_id
    month_cnt = 1
    vip_mode = 1
    mb_cost = 15000

    body = struct.pack(
        REQ_BODY_FMT,
        dest_id,
        account_id,
        month_cnt,
        vip_mode,
        mb_cost)
    pre_vfy = 'channelId=%u&securityCode=%s&data=' % (CHANNEL_ID, CHANNEL_KEY)
    vfy_code = hashlib.md5(pre_vfy + body).hexdigest()
    vfy_head = struct.pack(VERIFY_HEAD_FMT, CHANNEL_ID, vfy_code)

    pkg_len = PROTO_HEADER_SIZE + VERIFY_HEAD_SIZE + REQ_BODY_SIZE
    header = struct.pack(
        PROTO_HEADER_FMT,
        pkg_len,
        cmd_id,
        0,
        SEQ_NUM)

    return header + vfy_head + body


def send_request(host, port, cmd_id, account_id):
    packet = build_request(cmd_id, account_id)

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(TIMEOUT_SEC)
    try:
        sock.connect((host, port))
        sock.sendall(packet)
        header_prefix = recv_all(sock, 2)
        pkg_len = struct.unpack('<H', header_prefix)[0]
        if pkg_len < PROTO_HEADER_SIZE:
            raise RuntimeError('invalid pkg_len=%u' % pkg_len)
        payload = recv_all(sock, pkg_len - 2)
        return header_prefix + payload
    finally:
        sock.close()


def unpack_header(response):
    return struct.unpack(PROTO_HEADER_FMT, response[:PROTO_HEADER_SIZE])


def unpack_success_body(response):
    start = PROTO_HEADER_SIZE
    end = start + ACK_SUCCESS_BODY_SIZE
    return struct.unpack(ACK_SUCCESS_BODY_FMT, response[start:end])


def main():
    if len(sys.argv) != 2:
        print('usage: python2.7 test_mbvip_gate.py account_id')
        return 1

    account_id = int(sys.argv[1])

    try:
        response = send_request(HOST, PORT, CMD_ID, account_id)
    except Exception as err:
        print('cmd=%u' % CMD_ID)
        print('account_id=%u' % account_id)
        print('error=%s' % err)
        return 1

    pkg_len, cmd_id, result, seq_num = unpack_header(response)
    print('result=%u' % result)
    print('pkg_len=%u' % pkg_len)
    print('cmd_id=%u' % cmd_id)

    if result == 0 and len(response) >= PROTO_HEADER_SIZE + ACK_SUCCESS_BODY_SIZE:
        trans_id, mb_balance = unpack_success_body(response)
        print('trans_id=%u' % trans_id)
        print('mb_balance=%u' % mb_balance)

    return 0 if result == 0 else 1


if __name__ == '__main__':
    sys.exit(main())

