const url = require('url');
const axios = require('axios');
const qs = require('querystring')
const WebSocket = require('ws');

// disable certificate verification
process.env.NODE_TLS_REJECT_UNAUTHORIZED = '0';

const baseUrl = 'https://192.168.1.1';
const username = 'ubnt';
const password = 'Pa55w0rd';

(async () => {

    //console.error('Start');

    const sessionCookie = await logon(baseUrl, username, password);
    //console.error('Logged on');
    //console.error({ sessionCookie });

    await refreshHostNames(sessionCookie)

    const { protocol, hostname } = url.parse(baseUrl)
    const wsproto = protocol == 'https:' ? 'wss:' : 'ws:';

    const hostnameIsIp = (/\d+\.\d+\.\d+\.\d+/).test(hostname)
    const ws = new WebSocket(`${wsproto}//${hostname}/ws/stats`, {
        servername: hostnameIsIp ? '' : undefined // suppress "[DEP0123] DeprecationWarning: Setting the TLS ServerName to an IP address is not permitted by RFC 6066. This will be ignored in a future version."
    });

    ws.on('open', function open(x) {
        //console.error("Opening websocket");
        const initMessage = JSON.stringify({
            SUBSCRIBE: [{ name: "export" }],
            SESSION_ID: sessionCookie
        });

        //console.error("Sending: ", initMessage);
        ws.send(initMessage.length + '\n' + initMessage, function (e) {
            if (e)
                console.error('init message error', e);
        });

        //console.error('sent init message');
    });

    let messageLength = 0;
    let messageContent = '';

    ws.on('message', async (data) => {

        if (messageLength == 0) {
            //console.log('... new msg');
            const newlinepos = data.indexOf('\n');
            messageLength = ~~data.slice(0, newlinepos);
            messageContent = data.slice(newlinepos + 1);
        }
        else
        {
            //... append
            messageContent += data;
        }

        if (messageContent.length < messageLength) {
            // incomplete, wait for next part
            return;
        }

        await handleMessage(messageContent);

        try {
            await refreshHostNames(sessionCookie);
        } catch (error) {
            console.error('ERROR: refreshHostNames ', error);
        }

        messageLength = 0;
        messageContent = "";
    });

    ws.on('error', (code, reason) => {
        console.error('WS ERROR', { code, reason });
    })

    ws.on('close', (code, reason) => {
        console.error('WS CLOSE', { code, reason });
    })

})();

const _hostnames = {};
let _messageCount = 0

async function handleMessage(messageContent) {

    // header row
    if (_messageCount++ == 0) {
        console.log('time,categoryName,appName,hostname,ip,tx_bytes,tx_rate,rx_bytes,rx_rate');
    }

    /* {
        "export": {
            "192.168.1.66": {
                "Google Static Content(SSL)|Network protocols": {
                    "tx_bytes": "2355",
                    "tx_rate": "0",
                    "rx_bytes": "1500",
                    "rx_rate": "0"
                },
                "Youtube|Media streaming services": {
                    "tx_bytes": "85079",
                    "tx_rate": "0",
                    "rx_bytes": "688716",
                    "rx_rate": "0"
                },
            },
            ...    
    */
    //console.error(messageContent);
    const message = JSON.parse(messageContent);

    if (!message["export"])
        return;


    const time = new Date().toISOString().substring(0, 19).replace('T', ' ');

    const exportItems = message['export'];
    for (const ip in exportItems) {

        const exportItem = exportItems[ip];

        for (const appAndCategory in exportItem) {

            const stat = exportItem[appAndCategory]
            const hostname = _hostnames[ip] || ip;
            const [appName, categoryName] = appAndCategory.split(/\|/);

            console.log(`${time},${csvEncode(categoryName)},${csvEncode(appName)},${csvEncode(hostname)},${ip},${stat.tx_bytes},${stat.tx_rate},${stat.rx_bytes},${stat.rx_rate}`);
        }
    }
}


function csvEncode(str) {
    if (str == null)
        return '';
    if (str.indexOf(",") > -1 || str.indexOf("\"") > -1)
        str = '"' + str.replace(/"/g, '""') + '"';
    return str;
}

async function logon(baseUrl, username, password) {
    const form = { username: username, password: password };
    const response = await axios.post(baseUrl, qs.stringify(form), {
        maxRedirects: 0,
        validateStatus: () => true // accept all certs
    });

    if (!response.headers['set-cookie']){
        throw new Error('Logon failed, please check username/password')
    }
    
    const cookies = response.headers['set-cookie'].reduce((obj, item) => {
        const [name, value] = item.split(/=/);
        obj[name] = value.split(/;/)[0];
        return obj;
    }, {});
    //console.log({cookies});
    const sessionCookie = cookies['beaker.session.id'];
    return sessionCookie;
}

async function refreshHostNames(sessionCookie) {

    // only run every minute
    if ((new Date() - (refreshHostNames.lastUpdated || 0)) / 1000 < 60)
        return;

    //console.error('Refresh host names')
    const response = await axios.get(baseUrl.replace(/\/$/, '') + "/api/edge/data.json?data=dhcp_leases", {
        headers: { "Cookie": `beaker.session.id=${sessionCookie}` }
    });
    //console.log("refreshHostNames:", response.data.output['dhcp-server-leases'].LAN);

    const leases = response.data.output['dhcp-server-leases'].LAN;
    for (const ip in leases) {
        _hostnames[ip] = leases[ip]['client-hostname'].replace(/\?/g, '');
    }

    refreshHostNames.lastUpdated = new Date();
}