from scapy.packet import Packet
from scapy.fields import (
    EnumField,
    ShortField,
    StrLenField,
    BitEnumField,
    FieldLenField,
    StrFixedLenField,
)


class BJNPPkt(Packet):
    name = "BJNP Packet"

    BJNP_DEVICE_ENUM = {
        0x0: "Client",
        0x1: "Printer",
        0x2: "Scanner",
    }

    BJNP_COMMAND_ENUM = {
        0x000: "GetPortConfig",
        0x201: "GetNICInfo",
        0x202: "NICCmd",
        0x210: "SessionStart",
        0x211: "SessionEnd",
        0x212: "GetSessionInfo",
        0x220: "DataRead",
        0x221: "DataWrite",
        0x230: "GetDeviceID",
        0x232: "CmdNotify",
        0x240: "AppCmd",
    }

    BJNP_ERROR_ENUM = {
        0x8200: "Invalid header",
        0x8300: "Session error",
        0x8502: "Session already exists",
    }

    fields_desc = [
        StrFixedLenField("magic", default=b"MFNP", length=4),
        BitEnumField("device", default=0, size=1, enum=BJNP_DEVICE_ENUM),
        BitEnumField("cmd", default=0, size=15, enum=BJNP_COMMAND_ENUM),
        EnumField("err_no", default=0, enum=BJNP_ERROR_ENUM, fmt="!H"),
        ShortField("seq_no", default=0),
        ShortField("sess_id", default=0),
        FieldLenField("body_len", default=None, length_of="body", fmt="!I"),
        StrLenField("body", b"", length_from=lambda pkt: pkt.body_len),
    ]


def show_pkt(pkt, debug):
    if debug:
        print("\n--\n")
        pkt.show2()
        print(bytes(pkt))


def cleanup_payload(sock, debug=False):
    pkt = BJNPPkt(
        cmd=0x211,
        seq_no=0,
        sess_id=1,
    )
    show_pkt(pkt, debug)
    pkt = bytes(pkt)
    sock.sendall(pkt)

    res = BJNPPkt(sock.recv(4096))
    show_pkt(res, debug)


def store_payload(sock, bjnp_payload, debug=False):
    payload = bytes(bjnp_payload)

    if debug:
        print("\n--\n")
        print(bjnp_payload)
        print("")
        print(payload.hex())

    assert len(payload) <= 0x180, ValueError(
        "Payload too long: {} is greater than {}".format(len(payload), 0x180)
    )

    # Start a new session
    pkt = BJNPPkt(
        cmd=0x210,
        seq_no=0,
        sess_id=1,
        body=(b"\x00" * 8 + payload + b"\x00" * (0x180 - len(payload))),
    )
    show_pkt(pkt, debug)
    pkt = bytes(pkt)
    sock.sendall(pkt)

    res = BJNPPkt(sock.recv(4096))
    show_pkt(res, debug)

    # The printer should return a valid session ID
    assert res.sess_id != 0, ValueError("Failed to create session")
    sess_id = res.sess_id

    # Get the stored session info
    pkt = BJNPPkt(
        cmd=0x212,
        seq_no=1,
        sess_id=sess_id,
    )
    show_pkt(pkt, debug)
    pkt = bytes(pkt)
    sock.sendall(pkt)

    res = BJNPPkt(sock.recv(4096))
    show_pkt(res, debug)

    # Check that the payload was correctly stored
    assert res.body[4:].startswith(payload), ValueError(
        "Payload not found in session info"
    )
