#!/usr/bin/env python3
import argparse
import concurrent.futures
import csv
import ipaddress
import json
import socket
import struct
import sys
import time


def recv_until_timeout(sock, max_bytes=1048576):
    chunks = []
    total = 0

    while total < max_bytes:
        try:
            chunk = sock.recv(65536)
        except socket.timeout:
            break

        if not chunk:
            break

        chunks.append(chunk)
        total += len(chunk)

    return b"".join(chunks)


def recv_exact(sock, size):
    data = b""

    while len(data) < size:
        chunk = sock.recv(size - len(data))

        if not chunk:
            raise ConnectionError("connection closed")

        data += chunk

    return data


def tcp_open(ip, port, timeout):
    try:
        with socket.create_connection((str(ip), port), timeout=timeout):
            return True
    except Exception:
        return False


def legacy_command(ip, command, timeout):
    payload = json.dumps({"command": command}, separators=(",", ":")).encode("ascii")

    try:
        with socket.create_connection((str(ip), 4028), timeout=timeout) as sock:
            sock.settimeout(timeout)
            sock.sendall(payload)

            try:
                sock.shutdown(socket.SHUT_WR)
            except Exception:
                pass

            raw = recv_until_timeout(sock)
    except Exception as exc:
        return None, str(exc)

    text = raw.replace(b"\x00", b"").decode("utf-8", errors="replace").strip()

    if not text:
        return None, "empty response"

    try:
        return json.loads(text), None
    except Exception:
        return {"raw": text[:2000]}, None


def api3_get(ip, cmd, timeout):
    request = json.dumps({"cmd": cmd}, separators=(",", ":")).encode("ascii")
    frame = struct.pack("<I", len(request)) + request

    try:
        with socket.create_connection((str(ip), 4433), timeout=timeout) as sock:
            sock.settimeout(timeout)
            sock.sendall(frame)

            header = recv_exact(sock, 4)
            length = struct.unpack("<I", header)[0]

            if length < 1 or length > 10485760:
                return None, "bad api3 frame length"

            raw = recv_exact(sock, length)
    except Exception as exc:
        return None, str(exc)

    text = raw.decode("utf-8", errors="replace").strip()

    try:
        return json.loads(text), None
    except Exception:
        return {"raw": text[:2000]}, None


def walk_values(obj):
    if isinstance(obj, dict):
        for key, value in obj.items():
            yield key, value

            yield from walk_values(value)

    elif isinstance(obj, list):
        for item in obj:
            yield from walk_values(item)


def first_value(obj, names):
    wanted = {name.lower() for name in names}

    for key, value in walk_values(obj):
        if str(key).lower() in wanted and value not in (None, ""):
            return value

    return ""


def clean(value):
    if value is None:
        return ""

    text = str(value).replace("\n", " ").replace("\r", " ").strip()

    if len(text) > 60:
        return text[:57] + "..."

    return text


def scan_one(ip, timeout, query):
    result = {
        "ip": str(ip),
        "legacy_4028": False,
        "api3_4433": False,
        "mac": "",
        "model": "",
        "miner_sn": "",
        "power_sn": "",
        "hashrate": "",
        "legacy_error": "",
        "api3_error": "",
    }

    result["legacy_4028"] = tcp_open(ip, 4028, timeout)
    result["api3_4433"] = tcp_open(ip, 4433, timeout)

    if not query:
        return result

    legacy_data = None
    api3_data = None

    if result["legacy_4028"]:
        legacy_data, result["legacy_error"] = legacy_command(ip, "summary", timeout)

    if result["api3_4433"]:
        api3_data, result["api3_error"] = api3_get(ip, "get.device.info", timeout)

    combined = {
        "legacy": legacy_data or {},
        "api3": api3_data or {},
    }

    result["mac"] = clean(first_value(combined, [
        "mac", "mac_address", "net_mac", "ethaddr", "eth_mac"
    ]))

    result["model"] = clean(first_value(combined, [
        "miner_type", "miner type", "model", "type", "product_model"
    ]))

    result["miner_sn"] = clean(first_value(combined, [
        "miner_sn", "miner serial", "miner_sn_code", "serial", "sn"
    ]))

    result["power_sn"] = clean(first_value(combined, [
        "power_sn", "psu_sn", "power serial", "psu_serial"
    ]))

    result["hashrate"] = clean(first_value(combined, [
        "hashrate", "hash_rate", "mhs 5s", "mhs_5s", "ths 5s", "ths_5s", "total_hashrate"
    ]))

    return result


def build_targets(cidrs, ips, limit):
    targets = []

    for cidr in cidrs:
        network = ipaddress.ip_network(cidr, strict=False)

        for ip in network.hosts():
            targets.append(str(ip))

    for ip in ips:
        targets.append(str(ipaddress.ip_address(ip)))

    deduped = []
    seen = set()

    for ip in targets:
        if ip not in seen:
            seen.add(ip)
            deduped.append(ip)

    if limit and limit > 0:
        deduped = deduped[:limit]

    return deduped


def write_json(path, rows):
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(rows, handle, indent=2)


def write_csv(path, rows):
    fieldnames = [
        "ip",
        "legacy_4028",
        "api3_4433",
        "mac",
        "model",
        "miner_sn",
        "power_sn",
        "hashrate",
        "legacy_error",
        "api3_error",
    ]

    with open(path, "w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def print_table(rows, show_all):
    visible = rows if show_all else [
        row for row in rows if row["legacy_4028"] or row["api3_4433"]
    ]

    print("Scanned hosts:", len(rows))
    print("Hosts with miner ports open:", len(visible))
    print()

    if not visible:
        return

    columns = ["ip", "legacy_4028", "api3_4433", "mac", "model", "miner_sn", "power_sn", "hashrate"]
    widths = {
        "ip": 15,
        "legacy_4028": 11,
        "api3_4433": 10,
        "mac": 17,
        "model": 16,
        "miner_sn": 20,
        "power_sn": 20,
        "hashrate": 12,
    }

    header = "  ".join(col.ljust(widths[col]) for col in columns)
    print(header)
    print("-" * len(header))

    for row in visible:
        print("  ".join(clean(row.get(col, "")).ljust(widths[col])[:widths[col]] for col in columns))


def main():
    parser = argparse.ArgumentParser(description="Read-only Linux scanner for WhatsMiner ASICs.")
    parser.add_argument("--cidr", action="append", default=[], help="CIDR range to scan. Can be used multiple times.")
    parser.add_argument("--ip", action="append", default=[], help="Single IP to scan. Can be used multiple times.")
    parser.add_argument("--timeout", type=float, default=0.8, help="TCP timeout in seconds. Default: 0.8")
    parser.add_argument("--workers", type=int, default=256, help="Concurrent workers. Default: 256")
    parser.add_argument("--limit", type=int, default=0, help="Limit number of targets for testing.")
    parser.add_argument("--no-query", action="store_true", help="Only check open ports. Do not request miner info.")
    parser.add_argument("--show-all", action="store_true", help="Show hosts even when no miner ports are open.")
    parser.add_argument("--json", default="", help="Write JSON output to this file.")
    parser.add_argument("--csv", default="", help="Write CSV output to this file.")
    args = parser.parse_args()

    if not args.cidr and not args.ip:
        parser.error("Use --cidr and/or --ip")

    targets = build_targets(args.cidr, args.ip, args.limit)

    start = time.time()
    rows = []

    with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as executor:
        future_map = {
            executor.submit(scan_one, ip, args.timeout, not args.no_query): ip
            for ip in targets
        }

        for future in concurrent.futures.as_completed(future_map):
            rows.append(future.result())

    rows.sort(key=lambda row: tuple(int(part) for part in row["ip"].split(".")))

    print_table(rows, args.show_all)

    elapsed = time.time() - start
    print()
    print("Elapsed seconds:", round(elapsed, 2))

    if args.json:
        write_json(args.json, rows)
        print("JSON written:", args.json)

    if args.csv:
        write_csv(args.csv, rows)
        print("CSV written:", args.csv)


if __name__ == "__main__":
    main()