/*
 * stub_smbd.c — minimal SMB1 stub server for smbfs race reproduction (v2).
 *
 * Listens on 127.0.0.1:139 (NBSS), answers just enough NBSS+SMB1 for
 * DragonFly mount_smbfs to mount AND lookups to succeed, so smbnodes get
 * created in the hash table.
 *
 * Key protocol points:
 *  - NEGOTIATE: pick dialect index 7 ("NT LM 0.12"), WordCount=17
 *  - SESSION_SETUP_ANDX: minimal success (guest)
 *  - TREE_CONNECT_ANDX: success
 *  - TRANS2_FIND_FIRST2: return a fake file entry (any name works;
 *    smbfs uses the lookup name for the smbnode, not the found name)
 */
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>

static int verbose = 0;
#define V(...) do { if (verbose) fprintf(stderr, __VA_ARGS__); } while(0)

static unsigned char smb_sig[4] = {0xff,'S','M','B'};

static int readn(int fd, void *buf, int n) {
    int got = 0;
    while (got < n) {
        int r = read(fd, (char*)buf + got, n - got);
        if (r <= 0) return -1;
        got += r;
    }
    return 0;
}

/* Send a complete SMB1 response. words[] has nwords*2 bytes. bytes[] has nbytes. */
static void send_smb(int fd, unsigned char cmd, unsigned int mid,
                     unsigned int tid, unsigned int uid, unsigned int pid,
                     const unsigned char *words, int nwords,
                     const unsigned char *bytes, int nbytes)
{
    int hdr_sz = 32;
    int body_sz = 1 + nwords * 2 + 2 + nbytes;
    int total = hdr_sz + body_sz;
    int pkt_sz = 4 + total;
    unsigned char *pkt = calloc(1, pkt_sz + 256);
    int off;

    /* NBSS: type=0x00 + 3-byte big-endian length */
    pkt[0] = 0x00;
    pkt[1] = (total >> 16) & 0xff;
    pkt[2] = (total >> 8) & 0xff;
    pkt[3] = total & 0xff;
    off = 4;

    /* SMB header (32 bytes) */
    memcpy(pkt + off, smb_sig, 4); off += 4;     /* sig        */
    pkt[off++] = cmd;                             /* command    */
    pkt[off++] = 0; pkt[off++] = 0;              /* NT status  */
    pkt[off++] = 0; pkt[off++] = 0;
    pkt[off++] = 0x80;                            /* flags: reply */
    pkt[off++] = 0x00; pkt[off++] = 0x40;        /* flags2: NT_STATUS */
    pkt[off++] = (pid >> 8); pkt[off++] = (pid);  /* PIDHigh LE */
    memset(pkt + off, 0, 8); off += 8;            /* signature */
    pkt[off++] = 0; pkt[off++] = 0;              /* reserved */
    pkt[off++] = tid & 0xff; pkt[off++] = (tid >> 8);   /* TID LE */
    pkt[off++] = pid & 0xff; pkt[off++] = (pid >> 8);   /* PID LE */
    pkt[off++] = uid & 0xff; pkt[off++] = (uid >> 8);   /* UID LE */
    pkt[off++] = mid & 0xff; pkt[off++] = (mid >> 8);   /* MID LE */

    /* WordCount */
    pkt[off++] = nwords;
    if (words && nwords > 0)
        memcpy(pkt + off, words, nwords * 2);
    off += nwords * 2;
    /* ByteCount */
    pkt[off++] = nbytes & 0xff;
    pkt[off++] = (nbytes >> 8) & 0xff;
    if (bytes && nbytes > 0)
        memcpy(pkt + off, bytes, nbytes);

    write(fd, pkt, off);
    free(pkt);
}

/* Send a TRANS2 response with separate trans2-params and trans2-data sections.
 * The SMB TRANS2 response format:
 *   SMB header (32)
 *   WordCount = 10 (1)
 *   Words (10*2=20): TotalParamCount, TotalDataCount, Reserved,
 *                    ParamCount, ParamOffset, ParamDisp,
 *                    DataCount, DataOffset, DataDisp, Setup(2)
 *   ByteCount (2)
 *   [pad] [trans2 params] [trans2 data]
 */
static void send_trans2(int fd, unsigned int mid, unsigned int tid,
                        unsigned int uid, unsigned int pid,
                        const unsigned char *t2_params, int t2_paramlen,
                        const unsigned char *t2_data, int t2_datalen)
{
    int hdr_sz = 32;
    int wc_area = 1 + 10 * 2;   /* WordCount + 10 words */
    int bc_field = 2;
    /* trans2 params start right after ByteCount field, aligned to 2 or 4 */
    int pad1 = 0;  /* params can start immediately */
    int param_off = hdr_sz + wc_area + bc_field + pad1; /* offset from SMB header */
    int pad2 = (4 - ((param_off + t2_paramlen) % 4)) % 4; /* align data to 4 */
    if (t2_datalen == 0) pad2 = 0;
    int data_off = param_off + t2_paramlen + pad2;
    int byte_count = pad1 + t2_paramlen + pad2 + t2_datalen;
    int total = hdr_sz + wc_area + bc_field + byte_count;
    int pkt_sz = 4 + total;
    unsigned char *pkt = calloc(1, pkt_sz + 256);
    int off;

    /* NBSS */
    pkt[0] = 0x00;
    pkt[1] = (total >> 16) & 0xff;
    pkt[2] = (total >> 8) & 0xff;
    pkt[3] = total & 0xff;
    off = 4;

    /* SMB header */
    unsigned char cmd = 0x32; /* SMB_COM_TRANS2 */
    memcpy(pkt + off, smb_sig, 4); off += 4;
    pkt[off++] = cmd;
    pkt[off++] = 0; pkt[off++] = 0; pkt[off++] = 0; pkt[off++] = 0;
    pkt[off++] = 0x80; /* reply */
    pkt[off++] = 0x00; pkt[off++] = 0x40; /* NT_STATUS */
    pkt[off++] = (pid >> 8); pkt[off++] = pid;
    memset(pkt + off, 0, 8); off += 8;
    pkt[off++] = 0; pkt[off++] = 0;
    pkt[off++] = tid & 0xff; pkt[off++] = (tid >> 8);
    pkt[off++] = pid & 0xff; pkt[off++] = (pid >> 8);
    pkt[off++] = uid & 0xff; pkt[off++] = (uid >> 8);
    pkt[off++] = mid & 0xff; pkt[off++] = (mid >> 8);

    /* WordCount = 10 */
    pkt[off++] = 10;
    /* Words[0]: TotalParameterCount (LE) */
    pkt[off++] = t2_paramlen & 0xff; pkt[off++] = (t2_paramlen >> 8);
    /* Words[1]: TotalDataCount (LE) */
    pkt[off++] = t2_datalen & 0xff; pkt[off++] = (t2_datalen >> 8);
    /* Words[2]: Reserved */
    pkt[off++] = 0; pkt[off++] = 0;
    /* Words[3]: ParameterCount (LE) */
    pkt[off++] = t2_paramlen & 0xff; pkt[off++] = (t2_paramlen >> 8);
    /* Words[4]: ParameterOffset (LE) — from start of SMB header */
    pkt[off++] = param_off & 0xff; pkt[off++] = (param_off >> 8);
    /* Words[5]: ParameterDisplacement */
    pkt[off++] = 0; pkt[off++] = 0;
    /* Words[6]: DataCount (LE) */
    pkt[off++] = t2_datalen & 0xff; pkt[off++] = (t2_datalen >> 8);
    /* Words[7]: DataOffset (LE) — from start of SMB header */
    pkt[off++] = data_off & 0xff; pkt[off++] = (data_off >> 8);
    /* Words[8]: DataDisplacement */
    pkt[off++] = 0; pkt[off++] = 0;
    /* Words[9]: Setup (0 for FIND) */
    pkt[off++] = 0; pkt[off++] = 0;

    /* ByteCount */
    pkt[off++] = byte_count & 0xff;
    pkt[off++] = (byte_count >> 8) & 0xff;
    /* pad1 (0 bytes) */
    /* trans2 params */
    if (t2_params && t2_paramlen > 0)
        memcpy(pkt + off, t2_params, t2_paramlen);
    off += t2_paramlen;
    /* pad2 */
    off += pad2;
    /* trans2 data */
    if (t2_data && t2_datalen > 0)
        memcpy(pkt + off, t2_data, t2_datalen);
    off += t2_datalen;

    write(fd, pkt, off);
    free(pkt);
}

/* Build a FIND_FILE_DIRECTORY_INFO entry for a given filename.
 * Returns the entry size. Buffer must be large enough.
 * Format (SMB_FIND_FILE_DIRECTORY_INFO, info level 0x101):
 *   NextEntryOffset(4) FileIndex(4) CreationTime(8) LastAccessTime(8)
 *   LastWriteTime(8) ChangeTime(8) EndOfFile(8) AllocationSize(8)
 *   ExtFileAttributes(4) FileNameLength(4) FileName(FileNameLength bytes)
 */
static int build_find_entry(unsigned char *buf, const char *name, int is_dir) {
    int name_len = strlen(name);
    int entry_sz = 64 + name_len;  /* 64 fixed bytes + name (no null terminator) */
    int off = 0;
    /* NextEntryOffset = 0 (last/only entry) */
    buf[off++]=0; buf[off++]=0; buf[off++]=0; buf[off++]=0;
    /* FileIndex */
    buf[off++]=0; buf[off++]=0; buf[off++]=0; buf[off++]=0;
    /* CreationTime (8 bytes zero) */
    memset(buf+off, 0, 8); off += 8;
    /* LastAccessTime */
    memset(buf+off, 0, 8); off += 8;
    /* LastWriteTime */
    memset(buf+off, 0, 8); off += 8;
    /* ChangeTime */
    memset(buf+off, 0, 8); off += 8;
    /* EndOfFile (file size, 8 bytes LE) */
    memset(buf+off, 0, 8); off += 8;
    /* AllocationSize (8 bytes) */
    memset(buf+off, 0, 8); off += 8;
    /* ExtFileAttributes (4 bytes LE) */
    unsigned int attr = is_dir ? 0x10 : 0x20; /* DIR or ARCHIVE */
    buf[off++]=attr&0xff; buf[off++]=(attr>>8)&0xff;
    buf[off++]=(attr>>16)&0xff; buf[off++]=(attr>>24)&0xff;
    /* FileNameLength (4 bytes LE) — in bytes, NOT including null */
    buf[off++]=name_len&0xff; buf[off++]=(name_len>>8)&0xff;
    buf[off++]=(name_len>>16)&0xff; buf[off++]=(name_len>>24)&0xff;
    /* FileName — ASCII (not Unicode, since we said no UNICODE in flags2) */
    memcpy(buf+off, name, name_len); off += name_len;
    return off;
}

static void handle_session(int cfd) {
    unsigned char nbss[4];
    unsigned char buf[65536];
    unsigned int cur_tid = 1, cur_uid = 1;

    V("client connected\n");

    /* NBSS session request */
    if (readn(cfd, nbss, 4) < 0) return;
    if (nbss[0] == 0x81) {
        int nl = (nbss[1]<<16)|(nbss[2]<<8)|nbss[3];
        if (nl > 0 && nl < (int)sizeof(buf)) readn(cfd, buf, nl);
        unsigned char resp[4] = {0x82,0,0,0};
        write(cfd, resp, 4);
        V("NBSS positive response\n");
    }

    for (;;) {
        if (readn(cfd, nbss, 4) < 0) break;
        int msglen = (nbss[1]<<16)|(nbss[2]<<8)|nbss[3];
        if (nbss[0] != 0x00) break;
        if (msglen <= 0 || msglen >= (int)sizeof(buf)) break;
        if (readn(cfd, buf, msglen) < 0) break;
        if (memcmp(buf, smb_sig, 4) != 0) break;

        unsigned char cmd = buf[4];
        unsigned int mid = buf[30] | (buf[31] << 8);
        unsigned int tid = buf[24] | (buf[25] << 8);
        unsigned int pid_low = buf[26] | (buf[27] << 8);
        unsigned int pid_high = buf[12] | (buf[13] << 8);
        unsigned int pid = pid_low | (pid_high << 16);
        unsigned int uid = buf[28] | (buf[29] << 8);

        V("cmd=0x%02x mid=%u\n", cmd, mid);

        switch (cmd) {
        case 0x72: { /* NEGOTIATE */
            unsigned char w[17*2];
            memset(w, 0, sizeof(w));
            /* Word 0: DialectIndex = 7 ("NT LM 0.12") */
            w[0] = 7; w[1] = 0;
            /* Word 1: SecurityMode byte: 0x01 = user-level, no encrypt */
            w[2] = 0x01;
            /* Word 2: MaxMpxCount = 10 */
            w[4] = 10;
            /* Word 3: MaxVCs = 1 */
            w[6] = 1;
            /* Word 4-5: MaxBufferSize = 4356 (0x1104) — must be >= 4096 */
            w[8] = 0x04; w[9] = 0x11;
            /* Word 6-7: MaxRawSize = 1024 */
            w[12] = 0x00; w[13] = 0x04;
            /* Word 8-9: SessionKey */
            w[16] = 1;
            /* Word 10-11: Capabilities = NT_SMBS(0x10) | LARGE_FILES(0x08) */
            w[20] = 0x18; w[21] = 0;
            /* Word 12-15: SystemTime (8 bytes = 4 words) */
            /* Word 16: TimeZone (low byte) */
            /* Word 16 high byte: EncryptionKeyLength = 0 */
            send_smb(cfd, cmd, mid, tid, uid, pid, w, 17, NULL, 0);
            V("  NEGOTIATE -> dialect 7\n");
            break;
        }
        case 0x73: { /* SESSION_SETUP_ANDX */
            unsigned char w[3*2];
            memset(w, 0, sizeof(w));
            w[0] = 0xff; /* AndXCommand = none */
            cur_uid = uid;
            unsigned char d[] = {'D','F','B','S','D',0,'s','m','b','f','s',0,'W','O','R','K','G','R','O','U','P',0};
            send_smb(cfd, cmd, mid, tid, uid, pid, w, 3, d, sizeof(d));
            V("  SESSION_SETUP -> ok\n");
            break;
        }
        case 0x75: { /* TREE_CONNECT_ANDX */
            unsigned char w[7*2];
            memset(w, 0, sizeof(w));
            w[0] = 0xff;
            w[4] = 1; /* OptionalSupport */
            w[6] = 0xff; w[7] = 0x1f; /* MaxShareAccess */
            w[8] = 0xff; w[9] = 0x1f; /* GuestMaxShareAccess */
            cur_tid = tid;
            unsigned char d[] = {'A',':',0};
            send_smb(cfd, cmd, mid, tid, uid, pid, w, 7, d, sizeof(d));
            V("  TREE_CONNECT -> ok\n");
            break;
        }
        case 0x32: { /* TRANS2 */
            /* Parse the subcommand from setup words */
            int wc = buf[32];
            unsigned char *wp = buf + 33;
            int setup_count = 0;
            unsigned short subcmd = 0;
            if (wc >= 1) {
                /* For TRANS2, last word pair is the setup (subcommand) */
                /* SetupCount is at wp[wc*2-2], setup data at wp[wc*2-1..wc*2] */
                /* Actually: SetupCount is a separate field after all words in the byte area */
                /* For client->server TRANS2: WordCount, then WordCount words, then ByteCount,
                   then setup at specific position. Simplify: setup[0] is the last word pair. */
            }
            /* The subcommand is in the setup area. For SMB_COM_TRANS2, the
             * setup area contains 1 word: the subcommand.
             * Setup is located after the trans2 params/data in the byte area.
             * Setup[0] at: ByteCount field + 2 + SetupOffset... too complex.
             * Just respond to all TRANS2 as FIND_FIRST2. */
            subcmd = 0x0001; /* assume FIND_FIRST2 */

            if (subcmd == 0x0001 || subcmd == 0x0002) {
                /* FIND_FIRST2 or FIND_NEXT2 response */
                /* Trans2 params: SID(2) SearchCount(2) EndOfSearch(2) EaErrorOffset(2) LastNameOffset(2) */
                unsigned char t2p[10];
                memset(t2p, 0, sizeof(t2p));
                t2p[0] = 1; t2p[1] = 0; /* SID = 1 */
                t2p[2] = 1; t2p[3] = 0; /* SearchCount = 1 */
                t2p[4] = 1; t2p[5] = 0; /* EndOfSearch = 1 */
                /* Build a find entry for a fake file */
                unsigned char entry[256];
                int entry_sz = build_find_entry(entry, "FILE", 0);
                send_trans2(cfd, mid, tid, uid, pid, t2p, 10, entry, entry_sz);
                V("  TRANS2_FIND -> 1 entry \"FILE\"\n");
            } else {
                /* Other TRANS2: empty success */
                unsigned char t2p[2] = {0,0};
                send_trans2(cfd, mid, tid, uid, pid, t2p, 2, NULL, 0);
                V("  TRANS2 -> empty\n");
            }
            break;
        }
        case 0xa2: { /* NT_CREATE_ANDX */
            unsigned char w[42*2];
            memset(w, 0, sizeof(w));
            w[0] = 0xff; /* AndXCommand */
            w[4] = 0x01; /* FID = 1 */
            w[8] = 0x20; /* Disposition: open if exists */
            send_smb(cfd, cmd, mid, tid, uid, pid, w, 42, NULL, 0);
            V("  NT_CREATE -> FID=1\n");
            break;
        }
        case 0x04: /* CLOSE */
        case 0x71: /* TREE_DISCONNECT */
        case 0x74: /* LOGOFF_ANDX */
        case 0x2b: /* ECHO */
        case 0x6e: /* FIND_CLOSE2 */
            send_smb(cfd, cmd, mid, tid, uid, pid, NULL, 0, NULL, 0);
            V("  -> ok\n");
            break;
        default:
            send_smb(cfd, cmd, mid, tid, uid, pid, NULL, 0, NULL, 0);
            V("  -> default ok\n");
            break;
        }
    }
    V("session ended\n");
}

int main(int argc, char **argv) {
    int port = 139;
    int sfd, cfd, opt = 1;
    struct sockaddr_in addr;

    if (argc > 1) port = atoi(argv[1]);
    if (argc > 2) verbose = 1;

    signal(SIGPIPE, SIG_IGN);
    signal(SIGCHLD, SIG_IGN);

    sfd = socket(AF_INET, SOCK_STREAM, 0);
    setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = inet_addr("127.0.0.1");
    addr.sin_port = htons(port);
    if (bind(sfd, (struct sockaddr*)&addr, sizeof(addr)) < 0) { perror("bind"); return 1; }
    if (listen(sfd, 5) < 0) { perror("listen"); return 1; }
    fprintf(stderr, "stub_smbd on 127.0.0.1:%d\n", port);

    while ((cfd = accept(sfd, NULL, NULL)) >= 0) {
        if (fork() == 0) {
            close(sfd);
            handle_session(cfd);
            close(cfd);
            _exit(0);
        }
        close(cfd);
    }
    return 0;
}