#!/usr/bin/env node
/**
 * Tapo Device Discovery Tool
 * 
 * Discovers TP-Link Tapo devices on the local network using UDP broadcast.
 * Based on the protocol from the tapo-rs library.
 * 
 * Usage: node tapo-discover.js [broadcast_ip] [timeout_seconds]
 *   broadcast_ip: Default is 255.255.255.255
 *   timeout_seconds: Default is 10
 */

import dgram from 'dgram';
import crypto from 'crypto';
import { Buffer } from 'buffer';

// Tapo discovery port
const TAPO_DISCOVERY_PORT = 20002;
const DISCOVERY_INTERVAL_MS = 3000;

/**
 * Generate a simple RSA key pair for the discovery request
 * Tapo devices expect an RSA public key in PEM format
 */
function generateRsaKeyPair() {
    const { publicKey, privateKey } = crypto.generateKeyPairSync('rsa', {
        modulusLength: 1024,
        publicKeyEncoding: {
            type: 'pkcs1',
            format: 'pem'
        },
        privateKeyEncoding: {
            type: 'pkcs1',
            format: 'pem'
        }
    });
    return { publicKey, privateKey };
}

/**
 * CRC32 implementation (same polynomial as crc32fast in Rust)
 */
function crc32(data) {
    let crc = 0xFFFFFFFF;
    const table = getCrc32Table();

    for (let i = 0; i < data.length; i++) {
        crc = (crc >>> 8) ^ table[(crc ^ data[i]) & 0xFF];
    }

    return (crc ^ 0xFFFFFFFF) >>> 0;
}

let crc32Table = null;
function getCrc32Table() {
    if (crc32Table) return crc32Table;

    crc32Table = new Uint32Array(256);
    for (let i = 0; i < 256; i++) {
        let c = i;
        for (let j = 0; j < 8; j++) {
            c = (c & 1) ? (0xEDB88320 ^ (c >>> 1)) : (c >>> 1);
        }
        crc32Table[i] = c;
    }
    return crc32Table;
}

/**
 * Generate the AES discovery query packet
 * Based on tapo-rs: aes_discovery_query_generator.rs
 */
function generateDiscoveryPacket() {
    const { publicKey } = generateRsaKeyPair();

    // Create the JSON payload with RSA public key
    const keyPayload = JSON.stringify({
        params: {
            rsa_key: publicKey
        }
    });
    const keyPayloadBytes = Buffer.from(keyPayload, 'utf8');

    // Generate random device serial (4 bytes)
    const secret = crypto.randomBytes(4);
    const deviceSerial = secret.readUInt32BE(0);

    // Header fields (based on tapo-rs implementation)
    const version = 2;        // u8
    const msgType = 0;        // u8
    const opCode = 1;         // u16 BE
    const msgSize = keyPayloadBytes.length;  // u16 BE
    const flags = 17;         // u8
    const paddingByte = 0;    // u8
    const initialCrc = 0x5A6B7C8D;  // i32 BE (placeholder, will be replaced)

    // Build header (16 bytes total)
    const header = Buffer.alloc(16);
    let offset = 0;
    header.writeUInt8(version, offset++);           // 1 byte
    header.writeUInt8(msgType, offset++);           // 1 byte
    header.writeUInt16BE(opCode, offset);           // 2 bytes
    offset += 2;
    header.writeUInt16BE(msgSize, offset);          // 2 bytes
    offset += 2;
    header.writeUInt8(flags, offset++);             // 1 byte
    header.writeUInt8(paddingByte, offset++);       // 1 byte
    header.writeUInt32BE(deviceSerial, offset);     // 4 bytes
    offset += 4;
    header.writeInt32BE(initialCrc, offset);        // 4 bytes (placeholder)

    // Combine header and payload
    const query = Buffer.concat([header, keyPayloadBytes]);

    // Calculate CRC32 of the entire packet and update bytes 12-16
    const crcValue = crc32(query);
    query.writeUInt32BE(crcValue, 12);

    return query;
}

/**
 * Parse discovery response from Tapo device
 */
function parseDiscoveryResponse(data, rinfo) {
    const result = {
        ip: rinfo.address,
        port: rinfo.port,
        rawSize: data.length
    };

    // Response has 16-byte header + JSON payload
    if (data.length > 16) {
        try {
            const jsonStr = data.slice(16).toString('utf8');
            // Find the end of JSON (some responses have trailing garbage)
            const jsonEnd = jsonStr.lastIndexOf('}') + 1;
            if (jsonEnd > 0) {
                const json = JSON.parse(jsonStr.slice(0, jsonEnd));
                result.data = json;

                // Extract common fields
                if (json.result) {
                    result.deviceId = json.result.device_id;
                    result.owner = json.result.owner;
                    result.deviceType = json.result.device_type;
                    result.deviceModel = json.result.device_model;
                    result.factoryDefault = json.result.factory_default;
                    result.mgtEncryptSchm = json.result.mgt_encrypt_schm;
                }
            }
        } catch (e) {
            result.parseError = e.message;
            result.rawHex = data.toString('hex');
        }
    } else {
        result.rawHex = data.toString('hex');
    }

    return result;
}

/**
 * Main discovery function
 */
async function discoverDevices(broadcastAddr = '255.255.255.255', timeoutSeconds = 10) {
    return new Promise((resolve, reject) => {
        const devices = new Map();
        const socket = dgram.createSocket('udp4');

        socket.on('error', (err) => {
            console.error('Socket error:', err);
            socket.close();
            reject(err);
        });

        socket.on('message', (msg, rinfo) => {
            if (!devices.has(rinfo.address)) {
                const device = parseDiscoveryResponse(msg, rinfo);
                devices.set(rinfo.address, device);

                console.log(`\n✓ Found device at ${rinfo.address}`);
                if (device.deviceModel) {
                    console.log(`  Model: ${device.deviceModel}`);
                }
                if (device.deviceType) {
                    console.log(`  Type: ${device.deviceType}`);
                }
                if (device.deviceId) {
                    console.log(`  Device ID: ${device.deviceId}`);
                }
                if (device.mgtEncryptSchm) {
                    console.log(`  Encryption: ${JSON.stringify(device.mgtEncryptSchm)}`);
                }
                if (device.parseError) {
                    console.log(`  Parse error: ${device.parseError}`);
                }
            }
        });

        socket.bind(() => {
            socket.setBroadcast(true);

            console.log(`Discovering Tapo devices on ${broadcastAddr}:${TAPO_DISCOVERY_PORT} for ${timeoutSeconds} seconds...`);

            const sendDiscovery = () => {
                try {
                    const packet = generateDiscoveryPacket();
                    socket.send(packet, 0, packet.length, TAPO_DISCOVERY_PORT, broadcastAddr, (err) => {
                        if (err) {
                            console.error('Error sending discovery packet:', err.message);
                        }
                    });
                } catch (err) {
                    console.error('Error generating discovery packet:', err.message);
                }
            };

            // Send immediately and then at intervals
            sendDiscovery();
            const interval = setInterval(sendDiscovery, DISCOVERY_INTERVAL_MS);

            // Stop after timeout
            setTimeout(() => {
                clearInterval(interval);
                socket.close();
                resolve(Array.from(devices.values()));
            }, timeoutSeconds * 1000);
        });
    });
}

// Main execution
async function main() {
    const args = process.argv.slice(2).filter(a => !a.startsWith('--'));
    const broadcastAddr = args[0] || '255.255.255.255';
    const timeout = parseInt(args[1]) || 10;
    const jsonOutput = args.includes('--json');

    console.log('╔════════════════════════════════════════╗');
    console.log('║     Tapo Device Discovery Tool         ║');
    console.log('╚════════════════════════════════════════╝\n');

    const devices = await discoverDevices(broadcastAddr, timeout);

    console.log('\n────────────────────────────────────────');
    console.log(`Discovery complete. Found ${devices.length} device(s).`);

    if (devices.length > 0) {
        console.log('\nDiscovered Devices:');
        devices.forEach((device, idx) => {
            console.log(`  ${idx + 1}. ${device.ip} - ${device.deviceModel || 'Unknown Model'} (${device.deviceType || 'Unknown Type'})`);
        });
    } else {
        console.log('\nNo Tapo devices found on the network.');
        console.log('Tips:');
        console.log('  - Ensure devices are powered on and connected to the same network');
        console.log('  - Try specifying your network broadcast address (e.g., 192.168.3.255)');
        console.log('  - Some devices may not respond to UDP discovery');
    }

    if (jsonOutput && devices.length > 0) {
        console.log('\nJSON Output:');
        console.log(JSON.stringify(devices, null, 2));
    }
}

main().catch(console.error);