/*
 * Copyright (c) 2018-2020, Andreas Kling <kling@serenityos.org>
 * Copyright (c) 2021, Sergey Bugaev <bugaevc@serenityos.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include "Packet.h"
#include "Name.h"
#include "PacketHeader.h"
#include <AK/Debug.h>
#include <AK/MemoryStream.h>
#include <AK/StringBuilder.h>
#include <arpa/inet.h>

namespace DNS {

void Packet::add_question(Question const& question)
{
    m_questions.empend(question);

    VERIFY(m_questions.size() <= UINT16_MAX);
}

void Packet::add_answer(Answer const& answer)
{
    m_answers.empend(answer);

    VERIFY(m_answers.size() <= UINT16_MAX);
}

ErrorOr<ByteBuffer> Packet::to_byte_buffer() const
{
    PacketHeader header;
    header.set_id(m_id);
    if (is_query())
        header.set_is_query();
    else
        header.set_is_response();
    header.set_authoritative_answer(m_authoritative_answer);
    // FIXME: What should this be?
    header.set_opcode(0);
    header.set_response_code(m_code);
    header.set_truncated(false); // hopefully...
    header.set_recursion_desired(m_recursion_desired);
    // FIXME: what should the be for requests?
    header.set_recursion_available(m_recursion_available);
    header.set_question_count(m_questions.size());
    header.set_answer_count(m_answers.size());

    AllocatingMemoryStream stream;

    TRY(stream.write_value(header));
    for (auto& question : m_questions) {
        TRY(stream.write_value(question.name()));
        TRY(stream.write_value(htons((u16)question.record_type())));
        TRY(stream.write_value(htons(question.raw_class_code())));
    }
    for (auto& answer : m_answers) {
        TRY(stream.write_value(answer.name()));
        TRY(stream.write_value(htons((u16)answer.type())));
        TRY(stream.write_value(htons(answer.raw_class_code())));
        TRY(stream.write_value(htonl(answer.ttl())));
        if (answer.type() == RecordType::PTR) {
            Name name { answer.record_data() };
            TRY(stream.write_value(htons(name.serialized_size())));
            TRY(stream.write_value(name));
        } else {
            TRY(stream.write_value(htons(answer.record_data().length())));
            TRY(stream.write_until_depleted(answer.record_data().bytes()));
        }
    }

    auto buffer = TRY(ByteBuffer::create_uninitialized(stream.used_buffer_size()));
    TRY(stream.read_until_filled(buffer));
    return buffer;
}

class [[gnu::packed]] DNSRecordWithoutName {
public:
    DNSRecordWithoutName() = default;

    u16 type() const { return m_type; }
    u16 record_class() const { return m_class; }
    u32 ttl() const { return m_ttl; }
    u16 data_length() const { return m_data_length; }

    void* data() { return this + 1; }
    void const* data() const { return this + 1; }

private:
    NetworkOrdered<u16> m_type;
    NetworkOrdered<u16> m_class;
    NetworkOrdered<u32> m_ttl;
    NetworkOrdered<u16> m_data_length;
};

static_assert(sizeof(DNSRecordWithoutName) == 10);

ErrorOr<Packet> Packet::from_raw_packet(ReadonlyBytes bytes)
{
    if (bytes.size() < sizeof(PacketHeader)) {
        dbgln_if(LOOKUPSERVER_DEBUG, "DNS response not large enough ({} out of {}) to be a DNS packet", bytes.size(), sizeof(PacketHeader));
        return Error::from_string_literal("DNS response not large enough to be a DNS packet");
    }

    auto const& header = *bit_cast<PacketHeader const*>(bytes.data());
    dbgln_if(LOOKUPSERVER_DEBUG, "Got packet (ID: {})", header.id());
    dbgln_if(LOOKUPSERVER_DEBUG, "  Question count: {}", header.question_count());
    dbgln_if(LOOKUPSERVER_DEBUG, "    Answer count: {}", header.answer_count());
    dbgln_if(LOOKUPSERVER_DEBUG, " Authority count: {}", header.authority_count());
    dbgln_if(LOOKUPSERVER_DEBUG, "Additional count: {}", header.additional_count());

    Packet packet;
    packet.m_id = header.id();
    packet.m_query_or_response = header.is_response();
    packet.m_code = header.response_code();

    // FIXME: Should we parse further in this case?
    if (packet.code() != Code::NOERROR)
        return packet;

    size_t offset = sizeof(PacketHeader);

    for (u16 i = 0; i < header.question_count(); i++) {
        auto name = TRY(Name::parse(bytes, offset));
        struct RawDNSAnswerQuestion {
            NetworkOrdered<u16> record_type;
            NetworkOrdered<u16> class_code;
        };
        if (offset >= bytes.size() || bytes.size() - offset < sizeof(RawDNSAnswerQuestion))
            return Error::from_string_literal("Unexpected EOF when parsing DNS packet");

        auto const& record_and_class = *bit_cast<RawDNSAnswerQuestion const*>(bytes.offset_pointer(offset));
        u16 class_code = record_and_class.class_code & ~MDNS_WANTS_UNICAST_RESPONSE;
        bool mdns_wants_unicast_response = record_and_class.class_code & MDNS_WANTS_UNICAST_RESPONSE;
        packet.m_questions.empend(name, (RecordType)(u16)record_and_class.record_type, (RecordClass)class_code, mdns_wants_unicast_response);
        offset += 4;
        auto& question = packet.m_questions.last();
        dbgln_if(LOOKUPSERVER_DEBUG, "Question #{}: name=_{}_, type={}, class={}", i, question.name(), question.record_type(), question.class_code());
    }

    for (u16 i = 0; i < header.answer_count(); ++i) {
        auto name = TRY(Name::parse(bytes, offset));
        if (offset >= bytes.size() || bytes.size() - offset < sizeof(DNSRecordWithoutName))
            return Error::from_string_literal("Unexpected EOF when parsing DNS packet");

        auto const& record = *bit_cast<DNSRecordWithoutName const*>(bytes.offset_pointer(offset));
        offset += sizeof(DNSRecordWithoutName);
        if (record.data_length() > bytes.size() - offset)
            return Error::from_string_literal("Unexpected EOF when parsing DNS packet");

        ByteString data;

        switch ((RecordType)record.type()) {
        case RecordType::PTR: {
            size_t dummy_offset = offset;
            data = TRY(Name::parse(bytes, dummy_offset)).as_string();
            break;
        }
        case RecordType::CNAME:
            // Fall through
        case RecordType::A:
            // Fall through
        case RecordType::TXT:
            // Fall through
        case RecordType::AAAA:
            // Fall through
        case RecordType::SRV:
            data = ReadonlyBytes { record.data(), record.data_length() };
            break;
        default:
            // FIXME: Parse some other record types perhaps?
            dbgln("data=(unimplemented record type {})", (u16)record.type());
        }

        dbgln_if(LOOKUPSERVER_DEBUG, "Answer   #{}: name=_{}_, type={}, ttl={}, length={}, data=_{}_", i, name, record.type(), record.ttl(), record.data_length(), data);
        u16 class_code = record.record_class() & ~MDNS_CACHE_FLUSH;
        bool mdns_cache_flush = record.record_class() & MDNS_CACHE_FLUSH;
        packet.m_answers.empend(name, (RecordType)record.type(), (RecordClass)class_code, record.ttl(), data, mdns_cache_flush);
        offset += record.data_length();
    }

    return packet;
}

}