/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License (the "Licence").
 * You may not use this file except in compliance with the Licence.
 *
 * You can obtain a copy of the licence at RISC OS path @.^.LICENCE
 * or  http://www.riscosdev.com/lanman98/LICENCE.CDDL
 * See the Licence for the specific language governing permissions
 * and limitations under the Licence.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the Licence file. If applicable, add the
 * following below this CDDL HEADER, with the fields enclosed by
 * brackets "[]" replaced with your own identifying information:
 * Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */
 
/*
 *   Copyright 1996 Warm Silence Software Ltd.  All rights reserved.
 *   Use is subject to license terms.
 */

/*   PHBG 8/2/98: Initial version
 */

#include <ctype.h>
#include <string.h>
#include <stdio.h>
#include "LanMan98BaseLib/strext.h"
#include "LanMan98BaseLib/error.h"
#include "LanMan98BaseLib/memory.h"
#include "LanMan98BaseLib/blocked.h"
#include "LanMan98BaseLib/thread.h"
#include "word.h"
#include "LanMan98BaseLib/tcp.h"
#include "LanMan98BaseLib/time.h"
#include "var.h"
#include "discover.h"
#include "netbios.h"

// #define DIAGNOSTICS

#define RETRY_COUNT           (3)
#define UCAST_RETRY_TIMEOUT   (500)
#define BCAST_RETRY_TIMEOUT   (25)

#define RESOLVE_TIME_OUT      (2000)

/* For the naming service attempt that now precedes a broadcast.  Kept
   short because NetBIOS is still to be tried after it, and because a name
   that gets no answer here is looked up again, with the full timeout, when
   the connection is made. */
#define DNS_TIME_OUT          (300)
#define NAME_SERVICE_PORT     (137)
#define MAX_ADDRESS_SIZE      (256)
#define BUF_SIZE	      (2048)

#define MIN(x, y) ((x) < (y) ? (x) : (y))

typedef enum
{
    NAME_QUERY,
    NODE_STATUS
} query_id_t;

typedef struct
{
    int trid, flags,
        qdcount, ancount,
        nscount, arcount;
    char *qname;
    int qtype, qclass;
} query_t;

typedef struct
{
    int time_out;
    int naddr;
    char (*addr)[6];
} name_query_response_t;

typedef struct
{
    int time_out;
    int nname;
    char (*name)[18];
} node_status_response_t;

static void ensure_timeout(void)
{
    if(!ThreadTimedout())
    {
        ThreadResetTimeout();
        ExceptRethrow();
    }
}

static int query_poke(char *buf, query_t *nstat)
{
    char *p;

    p = buf;
    SetWord16BE(p, nstat->trid);              p += 2;
    SetWord16BE(p, nstat->flags);	      p += 2;
    SetWord16BE(p, nstat->qdcount);	      p += 2;
    SetWord16BE(p, nstat->ancount);	      p += 2;
    SetWord16BE(p, nstat->nscount);	      p += 2;
    SetWord16BE(p, nstat->arcount);	      p += 2;
    strcpy(p, nstat->qname);		      p += strlen(nstat->qname) + 1;
    SetWord16BE(p, nstat->qtype);	      p += 2;
    SetWord16BE(p, nstat->qclass);	      p += 2;
    return p - buf;
}

static int query_peek(char *buf, query_t *nstat)
{
    char *p;

    p = buf;
    nstat->trid = Word16BE(p);                p += 2;
    nstat->flags = Word16BE(p);	      	      p += 2;
    nstat->qdcount = Word16BE(p);	      p += 2;
    nstat->ancount = Word16BE(p);	      p += 2;
    nstat->nscount = Word16BE(p);	      p += 2;
    nstat->arcount = Word16BE(p);	      p += 2;
    nstat->qname = p;		              p += strlen(nstat->qname) + 1;
    /* query_poll terminates the datagram, so the name above ends inside it */
    nstat->qtype = Word16BE(p);	      	      p += 2;
    nstat->qclass = Word16BE(p);	      p += 2;
    return p - buf;
}

/*
    room is what is left of the datagram at buf.  The count of names is one
    byte, so a machine answering with 255 of them named four and a half
    kilobytes of names in a two kilobyte buffer, and every one of them was
    read - and one of them written to - outside it.
*/
static int node_status_response_peek(char *buf, int room, node_status_response_t *presp)
{
    char *p;
    int fits;

    p = buf;
    if(room < 7)
    {
        presp->nname = 0;
        presp->name = (char (*)[18]) p;
        return 0;
    }
    presp->time_out = Word32BE(p);                p += 6;
    presp->nname = *p++;
    presp->name = (char (*)[18]) p;
    fits = (room - 7) / 18;
    if(fits < 0)
        fits = 0;
    if(presp->nname > fits)
        presp->nname = fits;
    p += presp->nname * 18;
    return p - buf;
}

/*
    room is what is left of the datagram at buf: the count of addresses is
    the server's word, and six bytes of each of them are read.
*/
static int name_query_response_peek(char *buf, int room, name_query_response_t *presp)
{
    char *p;

    p = buf;
    if(room < 6)
    {
        presp->naddr = 0;
        presp->addr = (char (*)[6]) p;
        return 0;
    }
    presp->time_out = Word32BE(p);                p += 4;
    presp->naddr = Word16BE(p);	      	      	  p += 2;
    if(presp->naddr % 6 != 0)
        Error("Currupt address information from name query");
    if(presp->naddr > room - 6)
        presp->naddr = ((room - 6) / 6) * 6;
    presp->addr = (char (*)[6]) p;	      			  p += presp->naddr;
    presp->naddr /= 6;
    return p - buf;
}

/*
static void out_buf(char *buf, int i)
{
    FILE *f;

    f = fopen("CrudOut", "w");
    fwrite(buf, 1, i, f);
    fclose(f);
}

static void in_buf(char *buf, int i)
{
    FILE *f;

    f = fopen("CrudIn", "w");
    fwrite(buf, 1, i, f);
    fclose(f);
}
*/

static int encode_name(char *out, char *name)
{
    char buf[16], *p, *scope, *scope_buf;
    int i, len, j, c;

    scope_buf = VarRead("LanMan98$Scope");
    scope = scope_buf ? scope_buf : "";

    p = strchr(name, '.');
    len = p ? p - name : strlen(name);
    memset(buf, *name == '*' ? 0 : ' ', 16);
    memcpy(buf, name, MIN(len, 16));
    j = 0;
    out[j++] = 32;
    for(i = 0; i < 16; i++)
    {
        c = toupper(buf[i]);
        out[j++] = (c >> 4) + 'A';
        out[j++] = (c & 0xF) + 'A';
    }
    if(strlen(scope) + 34 > 255)
    {
        Free(scope_buf);
        Error("Scope ID too long");
    }
    while(strlen(scope))
    {
        p = strchr(scope, '.');
        len = p ? p - scope : strlen(scope);     /* not p - name: another string */
        out[j++] = len;
        memcpy(out + j, scope, len);
        j += len;
        scope += len + (p ? 1 : 0);
    }
    out[j++] = 0;
    Free(scope_buf);
    return j;
}

typedef struct
{
    query_id_t qid;
    udp_conn_t conn_in, conn_out;
    char *name, *name_server;
    tcp_addr_t addr;
    int trid;
    char *result;
    char nbnbuf[256];
    char buf[BUF_SIZE];
} *nq_rec_t;

static int base_trid = -1;

static void query_fin(nq_rec_t nq)
{
    Free(nq->name);
    Free(nq->name_server);
    TcpAddrDestruct(nq->addr);
    UdpDestruct(nq->conn_in);
    UdpDestruct(nq->conn_out);
    Free(nq);
}

static nq_rec_t query_init(char *name, char *name_server, query_id_t qid)
{
    nq_rec_t nq;

    nq = Malloc(sizeof(*nq));
    memset(nq, 0, sizeof(*nq));
    ExceptTry
    {
        nq->qid = qid;
        nq->name = strdup(name);
        nq->name_server = strdup(name_server);
        if(nq->name_server)
        {
            ThreadSetTimeout(RESOLVE_TIME_OUT);
            ExceptTry
            {
                nq->addr = TcpResolve(nq->name_server);
            }
            ExceptCatch
            {
                ensure_timeout();
            }
            ThreadResetTimeout();
            if(nq->addr == NULL)
                Error("Can't locate name server (WINS)");
        }
        base_trid = ((base_trid == -1 ? Time() : base_trid + 1) & 0x7FFF);
        nq->trid = base_trid;
        ExceptTry
        {
            nq->conn_in = nq->name_server ? UdpCreate(NAME_SERVICE_PORT) : NULL;
        } ExceptCatch {}
        nq->conn_out = UdpCreate(TCP_PORT_ANY);
        UdpSetPeer(NAME_SERVICE_PORT, nq->addr, nq->conn_out);
    }
    ExceptCatch
    {
        query_fin(nq);
        ExceptRethrow();
    }
    return nq;
}

static void query_request(nq_rec_t nq)
{
    query_t nstat;
    int used;

    nstat.trid = nq->trid;
    nstat.flags = nq->name_server ? 0x100 : 0x110;
    nstat.qdcount = 1;
    nstat.ancount = 0;
    nstat.nscount = 0;
    nstat.arcount = 0;
    encode_name(nq->nbnbuf, nq->name);
    nstat.qname = nq->nbnbuf;
    switch(nq->qid)
    {
        case NAME_QUERY: nstat.qtype = 0x20; break;
        case NODE_STATUS: nstat.qtype = 0x21; break;
    }
    nstat.qclass = 1;
    used = query_poke(nq->buf, &nstat);
    // out_buf(nq->buf, used);
    UdpWrite(nq->buf, used, nq->conn_out);
}

static void parse_addr(nq_rec_t nq, int offset, int size)
{
    name_query_response_t presp;

    if(offset < 0 || offset > size)
        return;
    name_query_response_peek(nq->buf + offset, size - offset, &presp);
    if(presp.naddr >= 1)
    {
        nq->result = nq->buf;
        sprintf(nq->buf, "%d.%d.%d.%d", presp.addr[0][2], presp.addr[0][3],
                                        presp.addr[0][4], presp.addr[0][5]);
    }
}

static void parse_names(nq_rec_t nq, int offset, int size)
{
    node_status_response_t presp;
    char *p;
    int i;

    /*
        The header is walked with strlen, so a datagram with no zero in it
        can put the offset past the end of what arrived.
    */
    if(offset < 0 || offset > size)
        return;
    node_status_response_peek(nq->buf + offset, size - offset, &presp);
    for(i = 0; i < presp.nname; i++)
    {
        if(presp.name[i][15] == ' ')
        {
            /* the name is sixteen bytes; the space that ends it is in them */
            p = memchr(presp.name[i], ' ', 16);
            if(p == NULL)
                continue;
            *p = 0;
            nq->result = presp.name[i];
            return;
        }
    }
}


static int query_poll(nq_rec_t nq)
{
    query_t nstat;
    int used, size;

    size = 0;
    ThreadSetTimeout(0);
    if(nq->conn_in)
    {
        ExceptTry
        {
            size = UdpRead(nq->buf, BUF_SIZE, nq->conn_in);
        }
        ExceptCatch
        {
            ensure_timeout();
        }
    }
    if(size == 0)
    {
        ExceptTry
        {
            size = UdpRead(nq->buf, BUF_SIZE, nq->conn_out);
        }
        ExceptCatch
        {
            ensure_timeout();
        }
    }
    ThreadResetTimeout();
    if(size < 12)
        return NULL;
    /*
        Nothing guarantees a zero anywhere in a datagram, and the header is
        walked with strlen: without this, a reply with no zero in it is
        measured off the end of the buffer.
    */
    if(size >= BUF_SIZE)
        size = BUF_SIZE - 1;
    nq->buf[size] = 0;
    // in_buf(nq->buf, size);
    used = query_peek(nq->buf, &nstat);
    if(nstat.trid == nq->trid
       && (nstat.flags & 0xF800) == 0x8000
       && nstat.ancount == 1
       && nstat.qdcount == 0
       && strcmp(nstat.qname, nq->nbnbuf) == 0
       && nstat.qclass == 1)
    {
        switch(nq->qid)
        {
            case NAME_QUERY:
                switch(nstat.qtype)
                {
                    case 0x20:
                        parse_addr(nq, used, size);
                        return 1;
                    case 0xA:
                        return nq->name_server ? 1 : 0;
                    default:
                        return 0;
                }
                break;
            case NODE_STATUS:
                switch(nstat.qtype)
                {
                    case 0x21:
                        parse_names(nq, used, size);
                        return 1;
                    default:
                        return 0;
                }
        }
    }
    return 0;
}

static char *do_query(char *name, char *name_server, query_id_t qid)
{
    nq_rec_t nq;
    int count_down, timeout, going;
    char *res;

    nq = query_init(name, name_server, qid);
    ExceptTry
    {
        count_down = RETRY_COUNT;
        timeout = Time();
        going = 1;
        while(going)
        {
            if(Time() - timeout >= 0)
            {
                timeout = Time() + (name_server ? UCAST_RETRY_TIMEOUT
                                                : BCAST_RETRY_TIMEOUT);
                if(count_down--)
                    query_request(nq);
                else
                    going = 0;
            }
            if(query_poll(nq))
                going = 0;
        }
    }
    ExceptCatch
    {
#ifdef DIAGNOSTICS
        printf("Failed: %s\n", ExceptCaught()->errmess);
#endif
        query_fin(nq);
        ExceptRethrow();
    }
    res = strdup(nq->result);
    query_fin(nq);
    return res;
}


int NetBIOSEncodeName(char *out, char *name)
{
    return encode_name(out, name);
}

char *NetBIOSForceResolve(char *name)
{
    char *res, *name_server;

#ifdef DIAGNOSTICS
    printf("Resolving %s\n", name);
#endif
    res = do_query(name, NULL, NAME_QUERY);
    if(res == NULL)
    {
#ifdef DIAGNOSTICS
        printf("Broadcast of name resolution failed\n");
#endif
        name_server = VarRead("LanMan98$NameServer");
        if(name_server)
        {
#ifdef DIAGNOSTICS
            printf("Trying name server %s\n", name_server);
#endif
            ExceptTry
            {
                res = do_query(name, name_server, NAME_QUERY);
            }
            ExceptCatch
            {
                Free(name_server);
                ExceptRethrow();
            }
            Free(name_server);
        }
    }
#ifdef DIAGNOSTICS
    if(res)
        printf("%s resolved to %s\n", name, res);
#endif
    return res;
}

/*
    Whether a name could be a NetBIOS name at all.  NetBIOS names are at
    most 15 characters and have no internal structure, so a dot rules one
    out: "server.local" and "192.168.1.5" are a DNS name and an address
    respectively, and broadcasting for either asks the network a question
    that has no possible answer.
*/
/*
    How long to wait for a machine to answer to its own name.  An answer to
    this comes back in milliseconds where it comes back at all, so this is
    a bound on failure rather than a wait anybody sits through.
*/
#define MDNS_RESOLVE_TIME (200)

/* A name in the domain the announcements use, and no other */
static int is_local_name(const char *name)
{
    int n = (int) strlen(name);
    const char *tail = ".local";
    int t = 6, i;

    if(n <= t)
        return 0;
    for(i = 0; i < t; i++)
        if(tolower((unsigned char) name[n - t + i]) != tail[i])
            return 0;
    return 1;
}

static int could_be_netbios(char *name)
{
    int n;

    for(n = 0; name[n]; n++)
        if(name[n] == '.')
            return 0;
    return (n > 0) && (n <= 15);
}

/*
    Ask the naming service, returning the address in the same dotted form a
    NetBIOS query would have produced, or NULL.  Failures are swallowed
    deliberately: every one of them just means trying NetBIOS next.
*/
static char *dns_resolve(char *name)
{
    volatile tcp_addr_t addr;
    unsigned char quad[4];
    char buf[24];

    addr = NULL;
    ThreadSetTimeout(DNS_TIME_OUT);
    ExceptTry
    {
        addr = TcpResolve(name);
    }
    ExceptCatch
    {
        addr = NULL;
    }
    ThreadResetTimeout();
    if(addr == NULL)
        return NULL;
    TcpAddrToQuad(addr, (char *) quad);
    TcpAddrDestruct(addr);
    sprintf(buf, "%d.%d.%d.%d", quad[0], quad[1], quad[2], quad[3]);
#ifdef DIAGNOSTICS
    printf("%s resolved to %s by the naming service\n", name, buf);
#endif
    return strdup(buf);
}

int NetBIOSIsAddress(const char *s)
{
    int dots, digits;

    dots = 0;
    digits = 0;
    for(; *s; s++)
    {
        if(*s == '.')
        {
            if(digits == 0)
                return 0;
            dots++;
            digits = 0;
        }
        else if((*s >= '0') && (*s <= '9'))
        {
            digits++;
        }
        else
        {
            return 0;
        }
    }
    return (dots == 3) && (digits > 0);
}

/*
    Whether something at an address takes a connection on either SMB port.

    Resolving a name only says what one source believes the address is.  A
    router's name server will often answer a bare name with an address that
    is nobody's - an old lease, the other port of a machine with two, or the
    router itself - while the machine answers to its own name correctly on
    the local network.  So an answer is tried before it is believed.  A
    machine on the local network accepts or refuses at once; the wait only
    bounds one that is not there.
*/
#define ANSWER_TIME (100)

static int answers_smb(const char *dotted)
{
    static const tcp_port_t ports[2] = { 445, 139 };
    tcp_addr_t volatile addr;
    tcp_conn_t volatile conn;
    int volatile ok;
    int i;

    addr = NULL;
    ExceptTry
    {
        addr = TcpResolve((char *) dotted);
    }
    ExceptCatch
    {
        addr = NULL;
    }
    if(addr == NULL)
        return 0;

    ok = 0;
    for(i = 0; (i < 2) && !ok; i++)
    {
        conn = NULL;
        ThreadSetTimeout(ANSWER_TIME);
        ExceptTry
        {
            conn = TcpCall(TCP_PORT_ANY, ports[i], (tcp_addr_t) addr);
        }
        ExceptCatch
        {
            conn = NULL;
        }
        ThreadResetTimeout();
        if(conn)
        {
            TcpConnDestruct((tcp_conn_t) conn);
            ok = 1;
        }
    }
    TcpAddrDestruct((tcp_addr_t) addr);
    return ok;
}

/*
    Resolve a name the way a Windows machine or a NAS does: every way there
    is, keeping the answer a machine is actually at.

    The ways are asked one at a time and each answer is tried before the
    next is asked, so a name the naming service gets right costs one
    connection attempt and nothing more.  An answer nothing takes a
    connection on is kept only in case nothing better turns up.  It is
    still returned then, so that whatever connects next reports the machine
    as not answering, which is what it is.

    The order is the naming service, the announcements, then NetBIOS, as
    Samba has always done.  LanMan98$NetBIOSFirst asks NetBIOS first, for
    the network where a short name means one thing to NetBIOS and something
    else to DNS.  A name ending in .local is only ever the announcements'.
*/
#define RESOLVE_WAYS (4)

char *NetBIOSResolve(char *name)
{
    char *tried[RESOLVE_WAYS];
    char *res, *var;
    int netbios_first, local, ntried, way, i, seen;

    if(NetBIOSIsAddress(name))
        return strdup(name);

    var = VarRead("LanMan98$NetBIOSFirst");
    netbios_first = (var != NULL);
    if(var)
        Free(var);
    local = is_local_name(name);

    ntried = 0;
    for(way = 0; way < RESOLVE_WAYS; way++)
    {
        res = NULL;
        switch(way)
        {
            case 0:
                if(!local && netbios_first && could_be_netbios(name))
                    res = NetBIOSForceResolve(name);
                break;
            case 1:
                if(!local)
                    res = dns_resolve(name);
                break;
            case 2:
                res = DiscoverResolve(name, MDNS_RESOLVE_TIME);
                break;
            case 3:
                if(!local && !netbios_first && could_be_netbios(name))
                    res = NetBIOSForceResolve(name);
                break;
        }
        if(res == NULL)
            continue;

        seen = 0;
        for(i = 0; i < ntried; i++)
            if(strcmp(tried[i], res) == 0)
                seen = 1;
        if(seen)
        {
            Free(res);
            continue;
        }
        if(answers_smb(res))
        {
            for(i = 0; i < ntried; i++)
                Free(tried[i]);
#ifdef DIAGNOSTICS
            printf("%s resolved to %s, which answers\n", name, res);
#endif
            return res;
        }
        tried[ntried++] = res;
    }

    if(ntried > 0)
    {
        for(i = 1; i < ntried; i++)
            Free(tried[i]);
        return tried[0];
    }
#ifdef DIAGNOSTICS
    printf("Can't resolve %s by any means\n", name);
#endif
    return strdup(name);
}

char *NetBIOSName(char *name, char *addr)
{
    char *res;

#ifdef DIAGNOSTICS
    printf("Broadcasting node status\n");
#endif
    res = do_query(name, NULL, NODE_STATUS);
    if(res == NULL)
    {
#ifdef DIAGNOSTICS
        printf("Failed: sending node status directly to PDC\n");
#endif
        res = do_query(name, addr, NODE_STATUS);
    }
    return res;
}
