package grpc import ( "context" "encoding/json" "fmt" "log/slog" "net/http" "net/http/cookiejar" "net/url" "sync" "time" pb "schedule_converter/proto/runner" "schedule_converter/auth" "schedule_converter/captcha" "schedule_converter/client" "schedule_converter/config" "schedule_converter/models" "schedule_converter/parser" apperr "schedule_converter/pkg/errors" ) type TaskHandler struct { cancelledTasks sync.Map } func NewTaskHandler() *TaskHandler { return &TaskHandler{} } func (h *TaskHandler) Execute(task *pb.Task) *pb.TaskResult { log := slog.With("task_id", task.TaskId, "task_type", task.Type.String()) log.Info("开始执行任务") result := &pb.TaskResult{ TaskId: task.TaskId, CompletedAt: time.Now().UnixMilli(), } if h.IsCancelled(task.TaskId) { result.Status = pb.TaskStatus_TASK_STATUS_CANCELLED result.ErrorMessage = "任务已取消" log.Info("任务已取消") return result } var err error var data []byte switch task.Type { case pb.TaskType_TASK_TYPE_GET_SCHEDULE: data, err = h.getSchedule(context.Background(), task.Payload) case pb.TaskType_TASK_TYPE_LOGIN: data, err = h.login(context.Background(), task.Payload) default: result.Status = pb.TaskStatus_TASK_STATUS_FAILED result.ErrorMessage = fmt.Sprintf("未知的任务类型: %v", task.Type) log.Warn("未知的任务类型") return result } if err != nil { result.Status = pb.TaskStatus_TASK_STATUS_FAILED result.ErrorMessage = err.Error() log.Error("任务执行失败", "error", err) } else { result.Status = pb.TaskStatus_TASK_STATUS_SUCCESS result.Data = data log.Info("任务执行成功") } return result } func (h *TaskHandler) CancelTask(taskID string) { h.cancelledTasks.Store(taskID, true) slog.Info("任务已标记为取消", "task_id", taskID) } func (h *TaskHandler) IsCancelled(taskID string) bool { v, ok := h.cancelledTasks.Load(taskID) return ok && v.(bool) } func (h *TaskHandler) getSchedule(ctx context.Context, payload []byte) ([]byte, error) { var req pb.GetSchedulePayload if err := json.Unmarshal(payload, &req); err != nil { return nil, apperr.New("parse_payload", err, "解析任务载荷失败") } log := slog.With("username", req.Username, "semester", req.Semester) log.Info("获取课表任务") httpClient, err := createHTTPClient() if err != nil { return nil, err } cfg := config.Get() setTWFIDCookie(httpClient, cfg.BaseURL, "") courses, err := executeLoginAndFetchSchedule(ctx, httpClient, req.Username, req.Password, req.Semester) if err != nil { return nil, err } result := &pb.ScheduleResultData{ StudentId: req.Username, Semester: req.Semester, Courses: convertCourses(courses), } log.Info("课表获取成功", "course_count", len(courses)) return json.Marshal(result) } func (h *TaskHandler) login(ctx context.Context, payload []byte) ([]byte, error) { var req pb.LoginPayload if err := json.Unmarshal(payload, &req); err != nil { return nil, apperr.New("parse_payload", err, "解析任务载荷失败") } log := slog.With("username", req.Username) log.Info("登录任务") httpClient, err := createHTTPClient() if err != nil { return nil, err } cfg := config.Get() setTWFIDCookie(httpClient, cfg.BaseURL, req.Twfid) if err := executeLogin(ctx, httpClient, req.Username, req.Password); err != nil { return nil, err } log.Info("登录成功") return json.Marshal(map[string]any{ "success": true, "message": "登录成功", }) } func createHTTPClient() (*http.Client, error) { jar, err := cookiejar.New(nil) if err != nil { return nil, apperr.New("create_jar", err, "创建 cookie jar 失败") } return &http.Client{Jar: jar}, nil } func setTWFIDCookie(httpClient *http.Client, baseURL, twfid string) { cfg := config.Get() if twfid == "" { twfid = cfg.TWFID } u, _ := url.Parse(baseURL) httpClient.Jar.SetCookies(u, []*http.Cookie{ {Name: "TWFID", Value: twfid}, }) } func executeLogin(ctx context.Context, httpClient *http.Client, username, password string) error { log := slog.With("username", username) if err := client.GetInitialCookiesWithClient(httpClient); err != nil { return apperr.New("get_cookies", err, "获取初始 Cookie 失败") } cfg := config.Get() _ = cfg.Recognizer.Timeout for attempt := 1; attempt <= 3; attempt++ { select { case <-ctx.Done(): return ctx.Err() default: } log.Info("尝试登录", "attempt", attempt, "max_retries", 3) captchaCode, err := captcha.GetAndRecognize(httpClient) if err != nil { log.Warn("验证码获取失败", "error", err, "attempt", attempt) time.Sleep(3 * time.Second) continue } success, err := auth.LoginWithClient(httpClient, captchaCode, username, password) if err != nil { log.Warn("登录请求失败", "error", err, "attempt", attempt) time.Sleep(3 * time.Second) continue } if success { return nil } log.Warn("登录验证失败", "attempt", attempt) time.Sleep(3 * time.Second) } return apperr.ErrLoginFailed } func executeLoginAndFetchSchedule(ctx context.Context, httpClient *http.Client, username, password, semester string) ([]models.Course, error) { if err := executeLogin(ctx, httpClient, username, password); err != nil { return nil, err } cfg := config.Get() if semester == "" { semester = cfg.Semester } return parser.FetchScheduleWithClient(httpClient, semester) } func convertCourses(courses []models.Course) []*pb.Course { result := make([]*pb.Course, len(courses)) for i, c := range courses { schedule := make([]*pb.ScheduleTime, len(c.Schedule)) for j, s := range c.Schedule { schedule[j] = &pb.ScheduleTime{ Day: s.Day, Sections: convertToInt32(s.Section), Weeks: convertToInt32(s.Weeks), } } result[i] = &pb.Course{ Name: c.Name, Teacher: c.Teacher, Room: c.Room, Schedule: schedule, } } return result } func convertToInt32(arr []int) []int32 { result := make([]int32, len(arr)) for i, v := range arr { result[i] = int32(v) } return result }