#include "nmap.h"
#include "hloop.h"
#include "hstring.h"
#include "hsocket.h"
#include "netinet.h"

using namespace hv;

#define MAX_RECVFROM_TIMEOUT 5000 // ms
#define MAX_SENDTO_PERSOCKET 1024

typedef struct nmap_ctx_s {
 Nmap* nmap;
 int send_cnt;
 int recv_cnt;
 int up_cnt;
 int idle_cnt;
} nmap_ctx_t;

static void on_idle(hidle_t* idle) {
 hloop_t* loop = hevent_loop(idle);
 nmap_ctx_t* ctx = (nmap_ctx_t*)hloop_userdata(loop);
 ctx->idle_cnt++;
 if (ctx->idle_cnt == 1) {
 // try again?
 }
 hloop_stop(loop);
}

static void on_timer(htimer_t* timer) {
 hloop_t* loop = hevent_loop(timer);
 hloop_stop(loop);
}

static void on_recvfrom(hio_t* io, void* buf, int readbytes) {
 //printd("on_recv fd=%d readbytes=%d\n", hio_fd(io), readbytes);
 /*
 char localaddrstr[SOCKADDR_STRLEN] = {0};
 char peeraddrstr[SOCKADDR_STRLEN] = {0};
 printd("[%s] <=> [%s]\n",
 SOCKADDR_STR(hio_localaddr(io), localaddrstr),
 SOCKADDR_STR(hio_peeraddr(io), peeraddrstr));
 */
 hloop_t* loop = hevent_loop(io);
 nmap_ctx_t* ctx = (nmap_ctx_t*)hloop_userdata(loop);
 if (++ctx->recv_cnt == ctx->send_cnt) {
 //hloop_stop(loop);
 }
 Nmap* nmap = ctx->nmap;
 struct sockaddr_in* peeraddr = (struct sockaddr_in*)hio_peeraddr(io);
 auto iter = nmap->find(peeraddr->sin_addr.s_addr);
 if (iter != nmap->end()) {
 if (iter->second == 0) {
 iter->second = 1;
 if (++ctx->up_cnt == nmap->size()) {
 hloop_stop(loop);
 }
 }
 }
}

int nmap_discover(Nmap* nmap) {
 hloop_t* loop = hloop_new(0);
 uint64_t start_hrtime = hloop_now_hrtime(loop);

 nmap_ctx_t ctx;
 ctx.nmap = nmap;
 ctx.send_cnt = 0;
 ctx.recv_cnt = 0;
 ctx.up_cnt = 0;
 ctx.idle_cnt = 0;
 hloop_set_userdata(loop, &ctx);

 char recvbuf[128];
 // icmp
 char sendbuf[44]; // 20IP + 44ICMP = 64
 icmp_t* icmp_req = (icmp_t*)sendbuf;
 icmp_req->icmp_type = ICMP_ECHO;
 icmp_req->icmp_code = 0;
 icmp_req->icmp_id = getpid();
 for (int i = 0; i < sizeof(sendbuf) - sizeof(icmphdr_t); ++i) {
 icmp_req->icmp_data[i] = i;
 }
 struct sockaddr_in peeraddr;
 hio_t* io = NULL;
 for (auto iter = nmap->begin(); iter != nmap->end(); ++iter) {
 if (iter->second == 1) continue;
 if (ctx.send_cnt % MAX_SENDTO_PERSOCKET == 0) {
 // socket
 int sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
 if (sockfd < 0) {
 perror("socket");
 if (errno == EPERM) {
 fprintf(stderr, "please use root or sudo to create a raw socket.\n");
 }
 return -socket_errno();
 }
 nonblocking(sockfd);
 so_sndbuf(sockfd, 425984); // 416K

 io = hio_get(loop, sockfd);
 if (io == NULL) return -1;
 hio_set_type(io, HIO_TYPE_IP);
 struct sockaddr_in localaddr;
 socklen_t addrlen = sizeof(localaddr);
 memset(&localaddr, 0, addrlen);
 localaddr.sin_addr.s_addr = htonl(INADDR_ANY);
 hio_set_localaddr(io, (struct sockaddr*)&localaddr, addrlen);
 hrecvfrom(loop, sockfd, recvbuf, sizeof(recvbuf), on_recvfrom);
 }
 icmp_req->icmp_seq = iter->first;
 icmp_req->icmp_cksum = 0;
 icmp_req->icmp_cksum = checksum((uint8_t*)icmp_req, sizeof(sendbuf));
 socklen_t addrlen = sizeof(peeraddr);
 memset(&peeraddr, 0, addrlen);
 peeraddr.sin_family = AF_INET;
 peeraddr.sin_addr.s_addr = iter->first;
 hio_set_peeraddr(io, (struct sockaddr*)&peeraddr, addrlen);
 hsendto(loop, hio_fd(io), sendbuf, sizeof(sendbuf), NULL);
 ++ctx.send_cnt;
 }

 htimer_add(loop, on_timer, MAX_RECVFROM_TIMEOUT, 1);
 hidle_add(loop, on_idle, 3);

 hloop_run(loop);
 uint64_t end_hrtime = hloop_now_hrtime(loop);
 hloop_free(&loop);

 // print result
 char ip[INET_ADDRSTRLEN];
 auto iter = nmap->begin();
 while (iter != nmap->end()) {
 inet_ntop(AF_INET, (void*)&iter->first, ip, sizeof(ip));
 printd("%s\t is %s.\n", ip, iter->second == 0 ? "down" : "up");
 ++iter;
 }
 printd("Nmap done: %lu IP addresses (%d hosts up) scanned in %.2f seconds\n",
 (unsigned long)nmap->size(), ctx.up_cnt, (end_hrtime-start_hrtime)/1000000.0f);

 return ctx.up_cnt;
}

int segment_discover(const char* segment16, Nmap* nmap) {
 StringList strlist = split(segment16, '.');
 if (strlist.size() != 4) return -1;
 uint32_t addr = 0;
 uint8_t* p = (uint8_t*)&addr;
 p[0] = atoi(strlist[0].c_str());
 p[1] = atoi(strlist[1].c_str());
 p[3] = 1;
 printd("Nmap scan %u.%u.x.1...\n", p[0], p[1]);
 nmap->clear();
 for (int i = 0; i < 256; ++i) {
 p[2] = i;
 (*nmap)[addr] = 0;
 }
 return nmap_discover(nmap);
}

int host_discover(const char* segment24, Nmap* nmap) {
 StringList strlist = split(segment24, '.');
 if (strlist.size() != 4) return -1;
 uint32_t addr = 0;
 uint8_t* p = (uint8_t*)&addr;
 p[0] = atoi(strlist[0].c_str());
 p[1] = atoi(strlist[1].c_str());
 p[2] = atoi(strlist[2].c_str());
 printd("Nmap scan %u.%u.%u.x...\n", p[0], p[1], p[2]);
 // 0,255 reserved
 nmap->clear();
 for (int i = 1; i < 255; ++i) {
 p[3] = i;
 (*nmap)[addr] = 0;
 }
 return nmap_discover(nmap);
}