Initial commit

2026-03-13 20:43:44 +08:00
commit 17cdb7ad0d
19 changed files with 4744 additions and 0 deletions
--- a/grpc/client.go
+++ b/grpc/client.go
@@ -0,0 +1,205 @@
+package grpc
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log"
+	"sync"
+	"time"
+
+	pb "jwts/proto/runner"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+// Client gRPC 客户端
+type Client struct {
+	serverAddr string
+	runnerID   string
+	version    string
+	conn       *grpc.ClientConn
+	stream     pb.RunnerHub_ConnectClient
+	mu         sync.Mutex
+	handler    *TaskHandler
+
+	// 配置
+	heartbeatInterval time.Duration
+}
+
+// NewClient 创建新的 gRPC 客户端
+func NewClient(serverAddr, runnerID, version string) *Client {
+	return &Client{
+		serverAddr:        serverAddr,
+		runnerID:          runnerID,
+		version:           version,
+		heartbeatInterval: 30 * time.Second,
+	}
+}
+
+// SetHandler 设置任务处理器
+func (c *Client) SetHandler(handler *TaskHandler) {
+	c.handler = handler
+}
+
+// Connect 连接到 gRPC 服务器
+func (c *Client) Connect(ctx context.Context) error {
+	log.Printf("正在连接到 gRPC 服务器: %s", c.serverAddr)
+
+	conn, err := grpc.Dial(c.serverAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		return fmt.Errorf("连接失败: %w", err)
+	}
+	c.conn = conn
+
+	client := pb.NewRunnerHubClient(conn)
+	stream, err := client.Connect(ctx)
+	if err != nil {
+		return fmt.Errorf("创建流失败: %w", err)
+	}
+	c.stream = stream
+
+	// 发送注册请求
+	if err := c.register(); err != nil {
+		return err
+	}
+
+	log.Println("gRPC 连接成功")
+	return nil
+}
+
+// register 发送注册请求
+func (c *Client) register() error {
+	msg := &pb.StreamMessage{
+		Message: &pb.StreamMessage_Register{
+			Register: &pb.RegisterRequest{
+				RunnerId: c.runnerID,
+				Version:  c.version,
+				Capabilities: map[string]string{
+					"TASK_TYPE_GET_SCHEDULE": "true",
+					"TASK_TYPE_LOGIN":        "true",
+				},
+			},
+		},
+	}
+	return c.stream.Send(msg)
+}
+
+// Run 运行主循环，接收消息
+func (c *Client) Run(ctx context.Context) error {
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+			msg, err := c.stream.Recv()
+			if err == io.EOF {
+				log.Println("服务器关闭了连接")
+				return nil
+			}
+			if err != nil {
+				return fmt.Errorf("接收消息失败: %w", err)
+			}
+
+			c.handleMessage(msg)
+		}
+	}
+}
+
+// handleMessage 处理接收到的消息
+func (c *Client) handleMessage(msg *pb.StreamMessage) {
+	switch m := msg.Message.(type) {
+	case *pb.StreamMessage_RegisterResponse:
+		log.Printf("注册成功: %s, 会话ID: %s, 心跳间隔: %d秒",
+			m.RegisterResponse.Message,
+			m.RegisterResponse.SessionId,
+			m.RegisterResponse.HeartbeatInterval)
+		if m.RegisterResponse.HeartbeatInterval > 0 {
+			c.heartbeatInterval = time.Duration(m.RegisterResponse.HeartbeatInterval) * time.Second
+		}
+	case *pb.StreamMessage_HeartbeatAck:
+		// 心跳确认，可以用于计算延迟
+		log.Printf("心跳确认，延迟: %dms", m.HeartbeatAck.Latency)
+	case *pb.StreamMessage_Task:
+		log.Printf("收到任务: %s, 类型: %v", m.Task.TaskId, m.Task.Type)
+		go c.handleTask(m.Task)
+	case *pb.StreamMessage_TaskCancel:
+		log.Printf("收到取消任务请求: %s, 原因: %s", m.TaskCancel.TaskId, m.TaskCancel.Reason)
+		if c.handler != nil {
+			c.handler.CancelTask(m.TaskCancel.TaskId)
+		}
+	default:
+		log.Printf("收到未知消息类型: %T", msg.Message)
+	}
+}
+
+// handleTask 处理任务
+func (c *Client) handleTask(task *pb.Task) {
+	if c.handler == nil {
+		log.Println("错误: 未设置任务处理器")
+		return
+	}
+
+	result := c.handler.Execute(task)
+
+	c.mu.Lock()
+	err := c.stream.Send(&pb.StreamMessage{
+		Message: &pb.StreamMessage_Result{
+			Result: result,
+		},
+	})
+	c.mu.Unlock()
+
+	if err != nil {
+		log.Printf("发送任务结果失败: %v", err)
+	} else {
+		log.Printf("任务结果已发送: %s, 状态: %v", result.TaskId, result.Status)
+	}
+}
+
+// sendHeartbeat 发送心跳
+func (c *Client) sendHeartbeat() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.stream == nil {
+		return fmt.Errorf("流未连接")
+	}
+
+	return c.stream.Send(&pb.StreamMessage{
+		Message: &pb.StreamMessage_Heartbeat{
+			Heartbeat: &pb.Heartbeat{
+				Timestamp: time.Now().UnixMilli(),
+			},
+		},
+	})
+}
+
+// HeartbeatLoop 心跳循环
+func (c *Client) HeartbeatLoop(ctx context.Context, interval time.Duration) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			if err := c.sendHeartbeat(); err != nil {
+				log.Printf("发送心跳失败: %v", err)
+			}
+		}
+	}
+}
+
+// Close 关闭连接
+func (c *Client) Close() error {
+	if c.stream != nil {
+		c.stream.CloseSend()
+	}
+	if c.conn != nil {
+		return c.conn.Close()
+	}
+	return nil
+}
--- a/grpc/handler.go
+++ b/grpc/handler.go
@@ -0,0 +1,305 @@
+package grpc
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"net/http"
+	"net/http/cookiejar"
+	"net/url"
+	"strings"
+	"sync"
+	"time"
+
+	pb "jwts/proto/runner"
+
+	"jwts/auth"
+	"jwts/captcha"
+	"jwts/config"
+	"jwts/models"
+	"jwts/parser"
+)
+
+// TaskHandler 任务处理器
+type TaskHandler struct {
+	cancelledTasks map[string]bool
+	mu             sync.RWMutex
+}
+
+// NewTaskHandler 创建新的任务处理器
+func NewTaskHandler() *TaskHandler {
+	return &TaskHandler{
+		cancelledTasks: make(map[string]bool),
+	}
+}
+
+// Execute 执行任务
+func (h *TaskHandler) Execute(task *pb.Task) *pb.TaskResult {
+	result := &pb.TaskResult{
+		TaskId:      task.TaskId,
+		CompletedAt: time.Now().UnixMilli(),
+	}
+
+	// 检查任务是否已取消
+	if h.IsCancelled(task.TaskId) {
+		result.Status = pb.TaskStatus_TASK_STATUS_CANCELLED
+		result.ErrorMessage = "任务已取消"
+		return result
+	}
+
+	switch task.Type {
+	case pb.TaskType_TASK_TYPE_GET_SCHEDULE:
+		data, err := h.getSchedule(task.Payload)
+		if err != nil {
+			result.Status = pb.TaskStatus_TASK_STATUS_FAILED
+			result.ErrorMessage = err.Error()
+		} else {
+			result.Status = pb.TaskStatus_TASK_STATUS_SUCCESS
+			result.Data = data
+		}
+	case pb.TaskType_TASK_TYPE_LOGIN:
+		data, err := h.login(task.Payload)
+		if err != nil {
+			result.Status = pb.TaskStatus_TASK_STATUS_FAILED
+			result.ErrorMessage = err.Error()
+		} else {
+			result.Status = pb.TaskStatus_TASK_STATUS_SUCCESS
+			result.Data = data
+		}
+	default:
+		result.Status = pb.TaskStatus_TASK_STATUS_FAILED
+		result.ErrorMessage = fmt.Sprintf("未知的任务类型: %v", task.Type)
+	}
+
+	return result
+}
+
+// CancelTask 取消任务
+func (h *TaskHandler) CancelTask(taskID string) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	h.cancelledTasks[taskID] = true
+	log.Printf("任务已标记为取消: %s", taskID)
+}
+
+// IsCancelled 检查任务是否已取消
+func (h *TaskHandler) IsCancelled(taskID string) bool {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+	return h.cancelledTasks[taskID]
+}
+
+// getSchedule 获取课表
+func (h *TaskHandler) getSchedule(payload []byte) ([]byte, error) {
+	var req pb.GetSchedulePayload
+	if err := json.Unmarshal(payload, &req); err != nil {
+		return nil, fmt.Errorf("解析任务载荷失败: %w", err)
+	}
+
+	semester := strings.TrimSpace(req.Semester)
+	if semester == "" {
+		semester = config.Semester
+	}
+
+	log.Printf("获取课表: 用户=%s, 学期=%s", req.Username, semester)
+
+	// 创建新的 HTTP 客户端
+	jar, err := cookiejar.New(nil)
+	if err != nil {
+		return nil, fmt.Errorf("创建 cookie jar 失败: %w", err)
+	}
+	httpClient := &http.Client{
+		Jar: jar,
+	}
+
+	// 设置 TWFID cookie (使用默认配置)
+	setTWFIDCookie(httpClient, config.BaseURL, "")
+
+	// 步骤1: 获取初始 Cookie
+	if err := getInitialCookies(httpClient); err != nil {
+		return nil, fmt.Errorf("获取初始 Cookie 失败: %w", err)
+	}
+
+	// 重试登录
+	maxRetries := 3
+	var loginSuccess bool
+	for attempt := 1; attempt <= maxRetries; attempt++ {
+		log.Printf("尝试登录 (第%d次)...", attempt)
+
+		// 获取并识别验证码
+		captchaCode, err := captcha.GetAndRecognize(httpClient)
+		if err != nil {
+			log.Printf("    验证码获取失败: %v", err)
+			continue
+		}
+
+		// 登录
+		success, err := loginWithCredentials(httpClient, captchaCode, req.Username, req.Password)
+		if err != nil {
+			log.Printf("    登录失败: %v", err)
+			continue
+		}
+
+		if success {
+			loginSuccess = true
+			break
+		}
+
+		log.Println("    登录失败，3秒后重试...")
+		time.Sleep(3 * time.Second)
+	}
+
+	if !loginSuccess {
+		return nil, fmt.Errorf("登录失败，已重试 %d 次", maxRetries)
+	}
+
+	// 获取课表
+	courses, err := parser.FetchScheduleWithClient(httpClient, semester)
+	if err != nil {
+		return nil, fmt.Errorf("获取课表失败: %w", err)
+	}
+
+	// 构建结果
+	result := &pb.ScheduleResultData{
+		StudentId: req.Username,
+		Semester:  semester,
+		Courses:   convertCourses(courses),
+	}
+
+	return json.Marshal(result)
+}
+
+// login 登录验证
+func (h *TaskHandler) login(payload []byte) ([]byte, error) {
+	var req pb.LoginPayload
+	if err := json.Unmarshal(payload, &req); err != nil {
+		return nil, fmt.Errorf("解析任务载荷失败: %w", err)
+	}
+
+	log.Printf("登录验证: 用户=%s", req.Username)
+
+	// 创建新的 HTTP 客户端
+	jar, err := cookiejar.New(nil)
+	if err != nil {
+		return nil, fmt.Errorf("创建 cookie jar 失败: %w", err)
+	}
+	httpClient := &http.Client{
+		Jar: jar,
+	}
+
+	// 设置 TWFID cookie
+	setTWFIDCookie(httpClient, config.BaseURL, req.Twfid)
+
+	// 步骤1: 获取初始 Cookie
+	if err := getInitialCookies(httpClient); err != nil {
+		return nil, fmt.Errorf("获取初始 Cookie 失败: %w", err)
+	}
+
+	// 重试登录
+	maxRetries := 3
+	for attempt := 1; attempt <= maxRetries; attempt++ {
+		log.Printf("尝试登录 (第%d次)...", attempt)
+
+		// 获取并识别验证码
+		captchaCode, err := captcha.GetAndRecognize(httpClient)
+		if err != nil {
+			log.Printf("    验证码获取失败: %v", err)
+			continue
+		}
+
+		// 登录
+		success, err := loginWithCredentials(httpClient, captchaCode, req.Username, req.Password)
+		if err != nil {
+			log.Printf("    登录失败: %v", err)
+			continue
+		}
+
+		if success {
+			return json.Marshal(map[string]interface{}{
+				"success": true,
+				"message": "登录成功",
+			})
+		}
+
+		log.Println("    登录失败，3秒后重试...")
+		time.Sleep(3 * time.Second)
+	}
+
+	return nil, fmt.Errorf("登录失败，已重试 %d 次", maxRetries)
+}
+
+// getInitialCookies 获取初始 Cookie
+func getInitialCookies(httpClient *http.Client) error {
+	log.Println("[1] 访问登录页面获取初始 Cookie...")
+
+	req, err := http.NewRequest("GET", config.BaseURL+"/loginNOCAS", nil)
+	if err != nil {
+		return err
+	}
+	req.Header = getHeaders()
+
+	resp, err := httpClient.Do(req)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	log.Printf("    登录页面状态码: %d", resp.StatusCode)
+	return nil
+}
+
+// loginWithCredentials 使用凭据登录
+func loginWithCredentials(httpClient *http.Client, captchaCode, username, password string) (bool, error) {
+	return auth.LoginWithClient(httpClient, captchaCode, username, password)
+}
+
+// setTWFIDCookie 设置 TWFID Cookie
+func setTWFIDCookie(httpClient *http.Client, baseURL, twfid string) {
+	if twfid == "" {
+		twfid = config.TWFID
+	}
+	u, _ := url.Parse(baseURL)
+	httpClient.Jar.SetCookies(u, []*http.Cookie{
+		{Name: "TWFID", Value: twfid},
+	})
+}
+
+// getHeaders 获取请求头
+func getHeaders() http.Header {
+	return http.Header{
+		"User-Agent":      []string{"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"},
+		"Accept":          []string{"text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"},
+		"Accept-Language": []string{"zh-CN,zh;q=0.9,en;q=0.8"},
+	}
+}
+
+// convertCourses 转换课程格式
+func convertCourses(courses []models.Course) []*pb.Course {
+	var result []*pb.Course
+	for _, c := range courses {
+		var schedule []*pb.ScheduleTime
+		for _, s := range c.Schedule {
+			schedule = append(schedule, &pb.ScheduleTime{
+				Day:      s.Day,
+				Sections: convertToInt32(s.Section),
+				Weeks:    convertToInt32(s.Weeks),
+			})
+		}
+		result = append(result, &pb.Course{
+			Name:     c.Name,
+			Teacher:  c.Teacher,
+			Room:     c.Room,
+			Schedule: schedule,
+		})
+	}
+	return result
+}
+
+// convertToInt32 将 []int 转换为 []int32
+func convertToInt32(arr []int) []int32 {
+	result := make([]int32, len(arr))
+	for i, v := range arr {
+		result[i] = int32(v)
+	}
+	return result
+}