// See the file "COPYING" in the main distribution directory for copyright.

#pragma once

#include "zeek/analyzer/protocol/tcp/TCP.h"

namespace zeek::analyzer::netbios_ssn {
namespace detail {

enum NetbiosSSN_Opcode : uint8_t {
    NETBIOS_SSN_MSG = 0x0,
    NETBIOS_DGM_DIRECT_UNIQUE = 0x10,
    NETBIOS_DGM_DIRECT_GROUP = 0x11,
    NETBIOS_DGM_BROADCAST = 0x12,
    NETBIOS_DGM_ERROR = 0x13,
    NETBIOS_DGG_QUERY_REQ = 0x14,
    NETBIOS_DGM_POS_RESP = 0x15,
    NETBIOS_DGM_NEG_RESP = 0x16,
    NETBIOS_SSN_REQ = 0x81,
    NETBIOS_SSN_POS_RESP = 0x82,
    NETBIOS_SSN_NEG_RESP = 0x83,
    NETBIOS_SSN_RETARG_RESP = 0x84,
    NETBIOS_SSN_KEEP_ALIVE = 0x85,
};

//  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
// |      TYPE     |     FLAGS     |            LENGTH             |
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

struct NetbiosSSN_RawMsgHdr {
    NetbiosSSN_RawMsgHdr(const u_char*& data, int& len);

    uint8_t type;
    uint8_t flags;
    uint16_t length;
};

//  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
// |   MSG_TYPE    |     FLAGS     |           DGM_ID              |
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
// |                           SOURCE_IP                           |
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
// |          SOURCE_PORT          |          DGM_LENGTH           |
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
// |         PACKET_OFFSET         |                               |
// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+                               |

struct NetbiosDGM_RawMsgHdr {
    NetbiosDGM_RawMsgHdr(const u_char*& data, int& len);

    uint8_t type;
    uint8_t flags;
    uint16_t id;
    uint32_t srcip;
    uint16_t srcport;
    uint16_t length;
    uint16_t offset;
};

enum NetbiosSSN_State : uint8_t {
    NETBIOS_SSN_TYPE,   // looking for type field
    NETBIOS_SSN_FLAGS,  // looking for flag field
    NETBIOS_SSN_LEN_HI, // looking for high-order byte of length
    NETBIOS_SSN_LEN_LO, // looking for low-order byte of length
    NETBIOS_SSN_BUF,    // building up the message in the buffer
};

class NetbiosSSN_Interpreter {
public:
    explicit NetbiosSSN_Interpreter(analyzer::Analyzer* analyzer);

    void ParseMessage(unsigned int type, unsigned int flags, const u_char* data, int len, bool is_query);

    // Version used when data points to type/flags/length.
    void ParseMessageTCP(const u_char* data, int len, bool is_query);
    void ParseMessageUDP(const u_char* data, int len, bool is_query);

    void Timeout() {}

protected:
    void ParseSessionMsg(const u_char* data, int len, bool is_query);
    void ParseSessionReq(const u_char* data, int len, bool is_query);
    void ParseSessionPosResp(const u_char* data, int len, bool is_query);
    void ParseSessionNegResp(const u_char* data, int len, bool is_query);
    void ParseRetArgResp(const u_char* data, int len, bool is_query);
    void ParseKeepAlive(const u_char* data, int len, bool is_query);

    // Datagram parsing
    void ParseBroadcast(const u_char* data, int len, bool is_query);
    void ParseDatagram(const u_char* data, int len, bool is_query);

    void ParseSambaMsg(const u_char* data, int len, bool is_query);

    void Event(EventHandlerPtr event, const u_char* data, int len, int is_orig = -1);

    // Pass in name/length, returns in xname/xlen the converted
    // name/length.  Returns 0 on failure; xname may still be
    // allocated and hold partial results at that point.
    int ConvertName(const u_char* name, int name_len, u_char*& xname, int& xlen);

protected:
    analyzer::Analyzer* analyzer;
    // SMB_Session* smb_session;
};

} // namespace detail

// ### This should be merged with TCP_Contents_RPC, TCP_Contents_DNS.
class Contents_NetbiosSSN final : public analyzer::tcp::TCP_SupportAnalyzer {
public:
    Contents_NetbiosSSN(Connection* conn, bool orig, detail::NetbiosSSN_Interpreter* interp);
    ~Contents_NetbiosSSN() override;

    void Flush(); // process any partially-received data

    detail::NetbiosSSN_State State() const { return state; }

protected:
    void DeliverStream(int len, const u_char* data, bool orig) override;
    void ProcessChunk(int& len, const u_char*& data, bool orig);

    detail::NetbiosSSN_Interpreter* interp;

    unsigned int type;
    unsigned int flags;

    u_char* msg_buf;
    int buf_n;    // number of bytes in msg_buf
    int buf_len;  // size of msg_buf
    int msg_size; // expected size of message

    detail::NetbiosSSN_State state;
};

class NetbiosSSN_Analyzer final : public analyzer::tcp::TCP_ApplicationAnalyzer {
public:
    explicit NetbiosSSN_Analyzer(Connection* conn);
    ~NetbiosSSN_Analyzer() override;

    void Done() override;
    void DeliverPacket(int len, const u_char* data, bool orig, uint64_t seq, const IP_Hdr* ip, int caplen) override;

    static analyzer::Analyzer* Instantiate(Connection* conn) { return new NetbiosSSN_Analyzer(conn); }

protected:
    void ConnectionClosed(analyzer::tcp::TCP_Endpoint* endpoint, analyzer::tcp::TCP_Endpoint* peer,
                          bool gen_event) override;
    void EndpointEOF(bool is_orig) override;

    void ExpireTimer(double t);

    detail::NetbiosSSN_Interpreter* interp;
    // SMB_Session* smb_session;
    Contents_NetbiosSSN* orig_netbios;
    Contents_NetbiosSSN* resp_netbios;
    int did_session_done;
};

// FIXME: Doesn't really fit into new analyzer structure. What to do?
int IsReuse(double t, const u_char* pkt);

} // namespace zeek::analyzer::netbios_ssn
