tdEngine_mqtt_mock/app/wits_sender.py

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

from config import build_wits_sender_dependencies
from model import REQUIRED_TRANSMISSION_CHANNELS, WITS_CHANNEL_MAPPING, WitsData

logger = logging.getLogger(__name__)
BEGIN_MARK = "&&\r\n"
END_MARK = "!!\r\n"
RECORD_TERMINATOR = "*\r\n"
RECONNECT_DELAY = 3
FIELD_RULES = {
    "deptbitm": (0.0, 20000.0, float),
    "chkp": (0.0, 20000.0, float),
    "sppa": (0.0, 20000.0, float),
    "rpma": (0, 400, int),
    "torqa": (0.0, 100000.0, float),
    "hkla": (0.0, 2000.0, float),
    "blkpos": (0.0, 1000.0, float),
    "woba": (0.0, 2000.0, float),
}


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


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


def build_random_wits_data(device_code):
    ts_ms = int(time.time() * 1000)
    hook_load = rand_float(17.3, 18.8)
    standpipe_pressure = rand_float(990.0, 1012.0)
    casing_pressure = rand_float(180.0, 260.0)
    rotary_rpm = rand_int(95, 135)
    torque = rand_float(8.0, 16.0)
    weight_on_bit = rand_float(6.0, 12.0)
    bit_depth = rand_float(199.8, 200.3)
    block_position = rand_float(5.8, 6.3)
    actcod = rand_int(1, 34)
    return WitsData(
        ts=ts_ms,
        wellid=device_code or "???1",
        stknum=0,
        recid=1,
        seqid=rand_int(1600, 9999),
        actual_date=time.strftime("%y%m%d"),
        actual_time=time.strftime("%H%M%S"),
        actual_ts=ts_ms,
        actcod=actcod,
        deptbitm=bit_depth,
        deptbitv=bit_depth - 1.45,
        deptmeas=bit_depth,
        deptvert=bit_depth - 1.45,
        blkpos=block_position,
        ropa=rand_float(0.8, 2.5),
        hkla=hook_load,
        hklx=hook_load,
        woba=weight_on_bit,
        wobx=-weight_on_bit,
        torqa=torque,
        torqx=torque,
        rpma=rotary_rpm,
        sppa=standpipe_pressure,
        chkp=casing_pressure,
        spm1=rand_int(98, 112),
        spm2=0,
        spm3=0,
        tvolact=rand_float(28.0, 31.0),
        tvolcact=rand_float(28.0, 31.0),
        mfop=0,
        mfoa=0.0,
        mfia=0.0,
        mdoa=rand_float(1069.8, 1070.1),
        mdia=26.846003,
        mtoa=29.113855,
        mtia=346.874634,
        mcoa=241.874634,
        mcia=0.0,
        stkc=0,
        lagstks=0,
        deptretm=bit_depth,
        gasa=0.0,
        space1=0.0,
        space2=0.0,
        space3=0.0,
        space4=0.0,
        space5=0.0,
    )


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


def validate_transmission_values(values):
    for field_name, (minimum, maximum, caster) in FIELD_RULES.items():
        raw_value = values.get(field_name)
        if raw_value is None or raw_value == "":
            raise ValueError(f"WITS field '{field_name}' is required")
        try:
            value = caster(raw_value)
        except (TypeError, ValueError) as exc:
            raise ValueError(f"WITS field '{field_name}' must be numeric, got {raw_value!r}") from exc
        if value < minimum or value > maximum:
            raise ValueError(
                f"WITS field '{field_name}' out of range [{minimum}, {maximum}], got {value}"
            )


def extract_channel_values(packet):
    lines = packet.replace("\r\n", "\n").replace("\r", "\n").split("\n")
    values = {}
    for raw_line in lines:
        line = raw_line.strip()
        if not line or line in {"&&", "!!", "*"}:
            continue
        if len(line) < 5:
            raise ValueError(f"Invalid WITS line: {line!r}")
        values[line[:4]] = line[4:]
    return values


def validate_packet(packet):
    channel_values = extract_channel_values(packet)
    missing_channels = [channel for channel in REQUIRED_TRANSMISSION_CHANNELS if channel not in channel_values]
    if missing_channels:
        missing_fields = [REQUIRED_TRANSMISSION_CHANNELS[channel] for channel in missing_channels]
        raise ValueError(f"WITS packet missing required fields: {', '.join(missing_fields)}")

    field_values = {
        field_name: channel_values[channel]
        for channel, field_name in REQUIRED_TRANSMISSION_CHANNELS.items()
    }
    validate_transmission_values(field_values)


def build_wits_packet(data):
    lines = [f"{channel}{format_wits_value(getattr(data, field_name), kind)}" for channel, field_name, kind in
             WITS_CHANNEL_MAPPING]
    packet = BEGIN_MARK + "\r\n".join(lines) + "\r\n" + END_MARK + RECORD_TERMINATOR
    validate_packet(packet)
    return packet


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]
    if lines and lines[-1] == "!!":
        lines = lines[:-1]
    packet = BEGIN_MARK + "\r\n".join(lines) + "\r\n" + END_MARK + RECORD_TERMINATOR
    validate_packet(packet)
    return packet


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


def open_connection(host, port, timeout):
    sock = socket.create_connection((host, port), timeout=timeout)
    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    sock.settimeout(timeout)
    return sock


def send_packet(sock, packet):
    sock.sendall(packet.encode("ascii", errors="strict"))


def run_wits_sender(args, deps):
    wits_config = deps.config.wits
    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
    interval = args.interval or 2.0

    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)",
        interval,
        args.count or "forever",
    )

    seq = 0
    sock = None
    try:
        while True:
            if sock is None:
                try:
                    sock = open_connection(host, port, timeout)
                    logger.info("WITS connected %s:%s", host, port)
                except ConnectionRefusedError:
                    logger.warning("WITS target refused connection %s:%s, retry in %ss", host, port, RECONNECT_DELAY)
                    time.sleep(RECONNECT_DELAY)
                    continue
                except TimeoutError:
                    logger.warning("WITS connect timeout %s:%s, retry in %ss", host, port, RECONNECT_DELAY)
                    time.sleep(RECONNECT_DELAY)
                    continue
                except OSError as exc:
                    logger.error("WITS connect failed %s:%s (%s), retry in %ss", host, port, exc, RECONNECT_DELAY)
                    time.sleep(RECONNECT_DELAY)
                    continue

            try:
                seq += 1
                if source_file:
                    packet = load_packet_from_file(source_file)
                else:
                    packet = build_wits_packet(build_random_wits_data(deps.config.tms.device_code))
                    # logging.info(f"packet: {packet}")
                send_packet(sock, 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(interval)
            except (BrokenPipeError, ConnectionResetError):
                logger.warning("WITS connection dropped by remote host, reconnecting in %ss", RECONNECT_DELAY)
                try:
                    sock.close()
                except OSError:
                    pass
                sock = None
                time.sleep(RECONNECT_DELAY)
            except TimeoutError:
                logger.warning("WITS send timeout, reconnecting in %ss", RECONNECT_DELAY)
                try:
                    sock.close()
                except OSError:
                    pass
                sock = None
                time.sleep(RECONNECT_DELAY)
            except OSError as exc:
                logger.error("WITS send failed (%s), reconnecting in %ss", exc, RECONNECT_DELAY)
                try:
                    sock.close()
                except OSError:
                    pass
                sock = None
                time.sleep(RECONNECT_DELAY)
    except KeyboardInterrupt:
        logger.info("WITS sender interrupted")
    finally:
        if sock is not None:
            try:
                sock.close()
            except OSError:
                pass
            logger.info("WITS disconnected")


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=2.0, help="Send interval in seconds")
    parser.add_argument("--count", type=int, default=0, 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()