用Go语言重写Linux系统命令 -- nc简化版

时间：2024-12-08 17:47:22

用Go语言重写Linux系统命令 – nc简化版

1. 引言

netcat，简称 nc，被誉为网络工具中的“瑞士军刀”，是网络调试与分析的利器。它的功能十分强大，然而平时我们经常使用的就是他的连通性测试功能，但是nc是被设计用来测试TCP和UDP的，对于HTTP支持不是太好，于是，我决定用 Go 语言重新实现一个简化版的 nc，专注于联通性测试功能，并扩展支持测试HTTP和Ping的连通性测试。通过这个项目，你不仅能深入学习网络编程，还能体验用 Go 语言开发系统工具的乐趣！

小目标：让我们的工具小而美，轻松搞定网络连通性调试！

2. 需求分析

在开始编写代码之前，我们需要明确工具的功能和设计目标：

  • 核心功能：

    1. TCP 测试：检查指定 IP 和端口是否可建立连接，支持数据发送与接收。
    2. UDP 测试：发送 UDP 数据包并接收响应。
    3. HTTP 测试：测试指定 URL 是否可访问，支持自定义请求方法（如 GET/POST）。
    4. Ping 测试：检查目标 IP 是否可达，并测量响应时间。

  • 用户体验设计：

    • 使用命令行参数指定测试类型和选项（如目标 IP、端口、超时）。
    • 错误提示清晰直观，方便排查问题。

  • 性能目标：快速响应，保持简洁。

3. 项目初始化

3.1 安装开发环境

确保已安装 Go 编译器（推荐版本 ≥ 1.20）。执行以下命令检查 Go 是否可用：

go version

3.2 初始化项目

创建项目目录并初始化 Go 模块：

mkdir gonc
cd gonc
go mod init gonc

项目结构：

gonc/
├── main.go        # 入口文件
├── tcp.go         # TCP 测试模块
├── udp.go         # UDP 测试模块
├── http.go        # HTTP 测试模块
├── ping.go        # Ping 测试模块
└── go.mod         # 依赖管理文件

4. 模块实现

4.1 实现 TCP 测试功能

在 tcp.go 中编写代码，建立 TCP 连接，发送和接收数据：

package mync

import (
	"fmt"
	"log"
	"net"
	"time"
)

// TestTCP 测试TCP端口是否连通，并支持发送和接收数据
func FireTCP(ip string, port, timeout int, data string) error {
	address := fmt.Sprintf("%s:%d", ip, port)

	// 设置连接超时时间
	conn, err := net.DialTimeout("tcp", address, time.Duration(timeout)*time.Second)
	if err != nil {
		log.Printf("Failed to connect to %s: %v\n", address, err)
		return err
	}
	defer conn.Close()

	fmt.Printf("Successfully connected to %s\n", address)

	// 发送数据
	if data != "" {
		fmt.Printf("Sending data: %s\n", data)
		_, err = conn.Write([]byte(data))
		if err != nil {
			log.Printf("Failed to send data: %v\n", err)
		} else {
			fmt.Println("Data sent successfully")
		}
	}

	// 接收响应数据
	buffer := make([]byte, 1024)
	conn.SetReadDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
	// 读取响应数据
	// 如果超时则提示用户, 服务器没有发送数据
	n, err := conn.Read(buffer)
	if err != nil {
		if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
			fmt.Println("No response from server")
		} else {
			log.Printf("Failed to read response: %v\n", err)
		}
	} else {
		fmt.Printf("Received: %s\n", string(buffer[:n]))
	}

	return err
}

4.2 实现 UDP 测试功能

在 udp.go 中实现 UDP 数据包的发送与接收：

package mync

import (
	"fmt"
	"log"
	"net"
	"time"
)

// TestUDP 测试UDP端口是否连通，并支持发送和接收数据
func FireUDP(ip string, port, timeout int, data string) error {
	address := fmt.Sprintf("%s:%d", ip, port)
	conn, err := net.Dial("udp", address)
	if err != nil {
		log.Printf("Failed to connect to %s: %v\n", address, err)
		return err
	}
	defer conn.Close()

	fmt.Printf("Successfully connected to %s\n", address)

	// 设置写超时
	conn.SetWriteDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
	_, err = conn.Write([]byte(data))
	if err != nil {
		log.Printf("Failed to send data: %v\n", err)
		return err
	}
	fmt.Println("Data sent successfully")

	// 设置读超时
	conn.SetReadDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
	buffer := make([]byte, 1024)
	n, err := conn.Read(buffer)
	if err != nil {
		log.Printf("No response received: %v\n", err)
	} else {
		fmt.Printf("Received: %s\n", string(buffer[:n]))
	}
	return err
}

4.3 实现 HTTP 测试功能

在 http.go 中构建 HTTP 请求：

package mync

import (
	"fmt"
	"io"
	"log"
	"net/http"
	"time"
)

// TestHTTP 测试HTTP/HTTPS连接
func FireHTTP(url, method string, headers map[string]string, timeout int) error {
	client := &http.Client{
		Timeout: time.Duration(timeout) * time.Second,
	}

	// 创建HTTP请求
	req, err := http.NewRequest(method, url, nil)
	if err != nil {
		log.Printf("Failed to create HTTP request: %v\n", err)
		return err
	}

	// 设置请求头
	for key, value := range headers {
		req.Header.Set(key, value)
	}

	// 发送HTTP请求
	resp, err := client.Do(req)
	if err != nil {
		log.Printf("HTTP request failed: %v\n", err)
		return err
	}
	defer resp.Body.Close()

	// 读取响应内容
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Printf("Failed to read response body: %v\n", err)
		return err
	}

	fmt.Printf("HTTP Status: %s\n", resp.Status)
	fmt.Printf("Response Body:\n%s\n", string(body))
	return err
}

4.4 实现 Ping 测试功能

在 ping.go 中：


import (
	"fmt"
	"log"
	"net"
	"os"
	"time"
)

const icmpEchoRequest = 8

// Ping 测试目标IP是否可达
func Ping(ip string, timeout int) error {
	conn, err := net.Dial("ip4:icmp", ip)
	if err != nil {
		log.Printf("Failed to connect to %s: %v\n", ip, err)
		return err
	}
	defer conn.Close()

	// 构造ICMP请求包
	id := os.Getpid() & 0xffff
	seq := 1
	msg := []byte{
		icmpEchoRequest, 0, 0, 0, // ICMP Type, Code, and Checksum placeholder
		byte(id >> 8), byte(id & 0xff), // Identifier
		byte(seq >> 8), byte(seq & 0xff), // Sequence number
	}
	checksum := calculateChecksum(msg)
	msg[2] = byte(checksum >> 8)
	msg[3] = byte(checksum & 0xff)

	// 设置写超时
	conn.SetDeadline(time.Now().Add(time.Duration(timeout) * time.Second))
	start := time.Now()
	_, err = conn.Write(msg)
	if err != nil {
		log.Printf("Failed to send ICMP request: %v\n", err)
		return err
	}

	// 接收ICMP响应
	buffer := make([]byte, 1024)
	_, err = conn.Read(buffer)
	if err != nil {
		log.Printf("Failed to receive ICMP response: %v\n", err)
		return err
	}

	elapsed := time.Since(start)
	fmt.Printf("Ping %s successful! Time=%v\n", ip, elapsed)
	return err
}

// calculateChecksum 计算ICMP校验和
func calculateChecksum(data []byte) uint16 {
	var sum int
	for i := 0; i < len(data)-1; i += 2 {
		sum += int(data[i])<<8 | int(data[i+1])
	}
	if len(data)%2 == 1 {
		sum += int(data[len(data)-1]) << 8
	}
	for (sum >> 16) > 0 {
		sum = (sum >> 16) + (sum & 0xffff)
	}
	return uint16(^sum)
}

4.5 命令行解析与功能整合

在 main.go 中，将上述模块整合，并解析命令行参数：

package main

import (
	"flag"
	"fmt"
	"net/url"
	"os"
	"strconv"
	"strings"

	"mync"
)

func main() {
	// 定义命令行参数
	var method, data string
	var httpHeaders httpHeaderMap
	var timeout int

	flag.StringVar(&method, "method", "GET", "HTTP method to use (default: GET)")
	flag.StringVar(&data, "data", "", "Data to send (for TCP or UDP)")
	flag.Var(&httpHeaders, "header", "Custom HTTP header (key=value pairs, can use multiple times)")
	flag.IntVar(&timeout, "timeout", 5, "Timeout in seconds (default: 5)")

	// 自定义 Usage
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s <protocol>://<target> [options]\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "options:\n")
		fmt.Fprintf(os.Stderr, "  -method <method>     HTTP method to use (default: GET)\n")
		fmt.Fprintf(os.Stderr, "  -data <data>         Data to send (for TCP or UDP)\n")
		fmt.Fprintf(os.Stderr, "  -header <header>   Custom HTTP header (key=value pairs, can use multiple times)\n")
		fmt.Fprintf(os.Stderr, "  -timeout <seconds>   Timeout in seconds (default: 5)\n")
		fmt.Fprintf(os.Stderr, "Examples:\n")
		fmt.Fprintf(os.Stderr, "  Test TCP connection:\n")
		fmt.Fprintf(os.Stderr, "    %s tcp://192.168.1.1:80 -data \"Hello, TCP!\" -timeout 10\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  Test UDP connection:\n")
		fmt.Fprintf(os.Stderr, "    %s udp://192.168.1.1:53 -data \"Hello, UDP!\" -timeout 5\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  Test HTTP connection:\n")
		fmt.Fprintf(os.Stderr, "    %s http://example.com -method GET -timeout 5\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  Ping a target:\n")
		fmt.Fprintf(os.Stderr, "    %s ping://192.168.1.1 -timeout 3\n", os.Args[0])
	}

	// 解析 flag 参数
	flag.Parse()

	// 必须至少有一个非 flag 参数
	args := flag.Args()
	if len(args) < 1 {
		fmt.Println("Error: Missing protocol and target. Example usage: tcp://192.168.1.1:80")
		flag.Usage()
		os.Exit(1)
	}

	// 解析第一个参数为协议和目标
	targetURL := args[0]
	parsedURL, err := url.Parse(targetURL)
	if err != nil || parsedURL.Scheme == "" || parsedURL.Host == "" {
		fmt.Printf("Error: Invalid protocol or target format: %s\n", targetURL)
		flag.Usage()
		os.Exit(1)
	}

	protocol := parsedURL.Scheme
	target := parsedURL.Host

	// 根据协议执行对应的功能
	switch protocol {
	case "tcp":
		ip, port, err := parseTarget(target)
		if err != nil {
			fmt.Println("Error:", err)
			os.Exit(1)
		}
		if err := mync.FireTCP(ip, port, timeout, data); err != nil {
			fmt.Println("TCP Test Error:", err)
		}

	case "udp":
		ip, port, err := parseTarget(target)
		if err != nil {
			fmt.Println("Error:", err)
			os.Exit(1)
		}
		if err := mync.FireUDP(ip, port, timeout, data); err != nil {
			fmt.Println("UDP Test Error:", err)
		}

	case "http", "https":
		headerMap := httpHeaders.ToMap()
		if err := mync.FireHTTP(targetURL, method, headerMap, timeout); err != nil {
			fmt.Println("HTTP Test Error:", err)
		}

	case "ping":
		if err := mync.Ping(target, timeout); err != nil {
			fmt.Println("Ping Test Error:", err)
		}

	default:
		fmt.Printf("Error: Unsupported protocol: %s\n", protocol)
		flag.Usage()
		os.Exit(1)
	}
}

// parseTarget 解析 <ip>:<port> 格式
func parseTarget(target string) (string, int, error) {
	parts := strings.Split(target, ":")
	if len(parts) != 2 {
		return "", 0, fmt.Errorf("invalid target format: %s, expected format <ip>:<port>", target)
	}
	ip := parts[0]
	port, err := strconv.Atoi(parts[1])
	if err != nil {
		return "", 0, fmt.Errorf("invalid port number: %s", parts[1])
	}
	return ip, port, nil
}

// httpHeaderMap 用于保存用户指定的多个 -header 参数
type httpHeaderMap []string

func (h *httpHeaderMap) String() string {
	return strings.Join(*h, ", ")
}

func (h *httpHeaderMap) Set(value string) error {
	*h = append(*h, value)
	return nil
}

// ToMap 将 httpHeaderMap 转换为 map
func (h *httpHeaderMap) ToMap() map[string]string {
	headerMap := make(map[string]string)
	for _, header := range *h {
		parts := strings.SplitN(header, "=", 2)
		if len(parts) == 2 {
			headerMap[parts[0]] = parts[1]
		}
	}
	return headerMap
}

5 测试

#!/bin/bash

set -e

ip=192.168.100.10
port=3333

timeout=2

# 测试tcp端口
echo ">>> Testing tcp port $port..."
./gonc tcp://$ip:$port -data "Hello" -timeout $timeout

# 测试udp端口
echo ">>> Testing udp port $port..."
./gonc udp://$ip:$port -data "Hello" -timeout $timeout

# 测试http服务
echo ">>> Testing http service..."
./gonc http://example.com -method GET -timeout $timeout

# 测试ping
echo ">>> Testing ping..."
sudo ./gonc ping://$ip -timeout $timeout

由于测试ping时使用了原始套接字,所以需要root权限

相关文章