tdEngine_mqtt_mock/app/wits_sender.py

import argparse
import logging
import random
import socket
import time
from datetime import datetime
from pathlib import Path

from config import build_wits_sender_dependencies
from model import WITS_FIELD_MAPPING, WitsData


logger = logging.getLogger(__name__)
BEGIN_MARK = "&&\r\n"
END_MARK = "!!\r\n"


def rand_int(a, b):
    return random.randint(a, b)


def rand_float(a, b, digits=2):
    return round(random.uniform(a, b), digits)


def build_random_wits_data(device_code):
    now = datetime.now()
    ts_ms = int(time.time() * 1000)
    return WitsData(
        ts=ts_ms,
        wellid=device_code,
        stknum=rand_int(0, 500),
        recid=0,
        seqid=rand_int(1, 999999),
        actual_date=float(now.strftime("%Y%m%d")),
        actual_time=float(now.strftime("%H%M%S")),
        actual_ts=ts_ms,
        actcod=rand_int(0, 9),
        actod_label="AUTO",
        deptbitm=rand_float(0, 5000),
        deptbitv=rand_float(0, 5000),
        deptmeas=rand_float(0, 5000),
        deptvert=rand_float(0, 5000),
        blkpos=rand_float(0, 100),
        ropa=rand_float(0, 200),
        hkla=rand_float(0, 500),
        hklx=rand_float(0, 500),
        woba=rand_float(0, 200),
        wobx=rand_float(0, 200),
        torqa=rand_float(0, 200),
        torqx=rand_float(0, 200),
        rpma=rand_int(0, 300),
        sppa=rand_float(0, 5000),
        chkp=rand_float(0, 5000),
        spm1=rand_int(0, 200),
        spm2=rand_int(0, 200),
        spm3=rand_int(0, 200),
        tvolact=rand_float(0, 20000),
        tvolcact=rand_float(0, 20000),
        mfop=rand_int(0, 1000),
        mfoa=rand_float(0, 1000),
        mfia=rand_float(0, 1000),
        mdoa=rand_float(0, 1000),
        mdia=rand_float(0, 1000),
        mtoa=rand_float(0, 1000),
        mtia=rand_float(0, 1000),
        mcoa=rand_float(0, 1000),
        mcia=rand_float(0, 1000),
        stkc=rand_int(0, 200),
        lagstks=rand_int(0, 200),
        deptretm=rand_float(0, 5000),
        gasa=rand_float(0, 100),
        space1=rand_float(0, 10),
        space2=rand_float(0, 10),
        space3=rand_float(0, 10),
        space4=rand_float(0, 10),
        space5=rand_float(0, 10),
    )


def format_wits_value(value, kind):
    if kind == "string":
        return str(value)
    if kind == "int":
        return str(int(value))
    return f"{float(value):.2f}"


def build_wits_packet(data):
    lines = []
    for index, field_name, kind in WITS_FIELD_MAPPING:
        value = getattr(data, field_name)
        lines.append(f"{index:02d}{format_wits_value(value, kind)}")
    return BEGIN_MARK + "\r\n".join(lines) + "\r\n" + END_MARK


def normalize_packet(text):
    body = text.replace("\r\n", "\n").replace("\r", "\n")
    lines = [line.rstrip() for line in body.split("\n") if line.strip()]
    if lines and lines[0] == "&&":
        lines = lines[1:]
    if lines and lines[-1] == "!!":
        lines = lines[:-1]
    return BEGIN_MARK + "\r\n".join(lines) + "\r\n" + END_MARK


def load_packet_from_file(path):
    return normalize_packet(Path(path).read_text(encoding="utf-8-sig"))


def send_packet(host, port, timeout, packet):
    with socket.create_connection((host, port), timeout=timeout) as sock:
        sock.sendall(packet.encode("ascii", errors="strict"))


def run_wits_sender(args, deps):
    wits_config = deps.config.wits
    device_code = deps.config.tms.device_code
    source_file = args.source_file or wits_config.source_file
    host = args.host or wits_config.host
    port = args.port or wits_config.port
    timeout = args.timeout or wits_config.timeout
    if not host or not port:
        raise ValueError("WITS target host/port is empty. Configure wits.host/wits.port or tms.server-ip/tms.server-port")

    logger.info("WITS sender config host=%s port=%s timeout=%ss source_file=%s interval=%ss count=%s", host, port, timeout, source_file or "(generated)", args.interval, args.count or "forever")

    seq = 0
    try:
        while True:
            seq += 1
            if source_file:
                packet = load_packet_from_file(source_file)
            else:
                packet = build_wits_packet(build_random_wits_data(device_code))
            send_packet(host, port, timeout, packet)
            logger.info("TX WITS #%s -> %s:%s", seq, host, port)
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug("WITS packet:\n%s", packet)
            if args.count and seq >= args.count:
                break
            time.sleep(args.interval)
    except KeyboardInterrupt:
        logger.info("WITS sender interrupted")


def add_arguments(parser):
    parser.add_argument("--config", default="config.yaml", help="Path to config yaml")
    parser.add_argument("--host", default="", help="Override target host")
    parser.add_argument("--port", type=int, default=0, help="Override target port")
    parser.add_argument("--timeout", type=int, default=0, help="Override socket timeout")
    parser.add_argument("--source-file", default="", help="Send raw WITS packet from file")
    parser.add_argument("--interval", type=float, default=3.0, help="Send interval in seconds")
    parser.add_argument("--count", type=int, default=1, help="Send count (0 = forever)")


def main(argv=None):
    parser = argparse.ArgumentParser(description="WITS TCP sender")
    add_arguments(parser)
    args = parser.parse_args(argv)
    deps = build_wits_sender_dependencies(args.config)
    run_wits_sender(args, deps)


if __name__ == "__main__":
    main()