testmontools/network_tools.py

from flask import Flask, request, jsonify
from urllib.parse import urlparse
from concurrent.futures import ThreadPoolExecutor
import subprocess
import re

app = Flask(__name__)

CURL_TIMING_FORMAT = (
    "dns_lookup=%(time_namelookup)s\n"
    "tcp_connect=%(time_connect)s\n"
    "tls_handshake=%(time_appconnect)s\n"
    "pre_transfer=%(time_pretransfer)s\n"
    "first_byte=%(time_starttransfer)s\n"
    "total=%(time_total)s\n"
    "size_bytes=%(size_download)s\n"
    "speed_bps=%(speed_download)s\n"
    "http_code=%(http_code)s\n"
).replace('%', '%')

CURL_NOISE = re.compile(r'[\{\}] \[\d+ bytes data\]|\[[\d\s]+ bytes data\]|\r')


def extract_host(value: str) -> str:
    if value.startswith(('http://', 'https://')):
        return urlparse(value).hostname or value
    return value


def parse_input(data: dict):
    raw = data.get('url') or data.get('host')
    if not raw:
        return None, None
    return raw, extract_host(raw)


def parse_curl_timing(text: str) -> dict:
    timing = {}
    for line in text.splitlines():
        if '=' in line:
            key, _, val = line.partition('=')
            try:
                timing[key.strip()] = float(val.strip())
            except ValueError:
                timing[key.strip()] = val.strip()
    return timing


def clean_curl_output(text: str) -> str:
    return '\n'.join(
        line for line in text.splitlines()
        if not CURL_NOISE.search(line)
    )


def run_ping(host: str) -> dict:
    result = subprocess.run(
        ['ping', '-c', '4', host],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
    )
    return {
        'resolved_host': host,
        'stdout': result.stdout,
        'stderr': result.stderr,
        'returncode': result.returncode
    }


def run_traceroute(host: str) -> dict:
    result = subprocess.run(
        ['traceroute', host],
        stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
    )
    return {
        'resolved_host': host,
        'stdout': result.stdout,
        'stderr': result.stderr,
        'returncode': result.returncode
    }



def run_curl(url: str, headers_only: bool = False) -> dict:
    args = [
        'curl', '-sS', '-v',
        '-w', '\n===TIMING===\n' + CURL_TIMING_FORMAT,
        '--max-time', '30',
    ]

    if headers_only:
        args.append('-I')

    args.append(url)

    result = subprocess.run(
        args,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )

    # Разделяем тело и timing
    parts = result.stdout.split('===TIMING===')

    body = parts[0].strip()
    timing_text = parts[1] if len(parts) > 1 else ""

    return {
        'url': url,
        'timing': parse_curl_timing(timing_text),
        'trace': clean_curl_output(result.stderr),
        'returncode': result.returncode,
        'body': body
    }

def run_dns(host: str, record_types=None) -> dict:
    if record_types is None:
        record_types = ['A', 'AAAA']

    results = {}
    all_ips = []

    for rtype in record_types:
        result = subprocess.run(
            ['dig', '+short', rtype, host],
            stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
        )
        records = [line.strip() for line in result.stdout.splitlines() if line.strip()]
        results[rtype] = records
        if rtype in ('A', 'AAAA'):
            all_ips.extend(records)

    return {
        'resolved_host': host,
        'records': results,
        'all_ips': all_ips
    }

# ── Endpoints ────────────────────────────────────────────────────────────────

@app.route('/testmontools/ping', methods=['POST'])
def ping():
    data = request.get_json()
    raw, host = parse_input(data or {})
    if not host:
        return jsonify({'error': 'Field "url" or "host" is required'}), 400
    return jsonify({'input': raw, **run_ping(host)})


@app.route('/testmontools/traceroute', methods=['POST'])
def traceroute():
    data = request.get_json()
    raw, host = parse_input(data or {})
    if not host:
        return jsonify({'error': 'Field "url" or "host" is required'}), 400
    return jsonify({'input': raw, **run_traceroute(host)})


@app.route('/testmontools/curl', methods=['POST'])
def curl():
    data = request.get_json()
    raw, _ = parse_input(data or {})
    if not raw:
        return jsonify({'error': 'Field "url" or "host" is required'}), 400
    url = raw if raw.startswith(('http://', 'https://')) else f'https://{raw}'
    return jsonify({'input': raw, **run_curl(url, headers_only=bool(data.get('headers')))})

@app.route('/testmontools/dns', methods=['POST'])
def dns():
    data = request.get_json()
    raw, host = parse_input(data or {})
    if not host:
        return jsonify({'error': 'Field "url" or "host" is required'}), 400

    record_types = data.get('types')  # optional, e.g. ["A", "AAAA", "MX", "TXT", "CNAME", "NS"]
    if record_types:
        record_types = [t.upper() for t in record_types]

    return jsonify({'input': raw, **run_dns(host, record_types)})
    
@app.route('/testmontools/all', methods=['POST'])
def check_all():
    data = request.get_json()
    raw, host = parse_input(data or {})
    if not raw or not host:
        return jsonify({'error': 'Field "url" or "host" is required'}), 400

    url = raw if raw.startswith(('http://', 'https://')) else f'https://{raw}'

    # Запускаем ping и traceroute параллельно, curl — отдельно
    with ThreadPoolExecutor(max_workers=3) as pool:
        f_ping       = pool.submit(run_ping, host)
        f_traceroute = pool.submit(run_traceroute, host)
        f_curl       = pool.submit(run_curl, url)

    return jsonify({
        'input': raw,
        'curl':       f_curl.result(),
        'ping':       f_ping.result(),
        'traceroute': f_traceroute.result(),
    })


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7088, debug=False)