mihomo/constant/rewrite.go

package constant

import (
	"encoding/json"
	"errors"
	regexp "github.com/dlclark/regexp2"
)

var RewriteTypeMapping = map[string]RewriteType{
	MitmReject.String():         MitmReject,
	MitmReject200.String():      MitmReject200,
	MitmRejectImg.String():      MitmRejectImg,
	MitmRejectDict.String():     MitmRejectDict,
	MitmRejectArray.String():    MitmRejectArray,
	Mitm302.String():            Mitm302,
	Mitm307.String():            Mitm307,
	MitmRequestHeader.String():  MitmRequestHeader,
	MitmRequestBody.String():    MitmRequestBody,
	MitmResponseHeader.String(): MitmResponseHeader,
	MitmResponseBody.String():   MitmResponseBody,
}

const (
	MitmReject RewriteType = iota
	MitmReject200
	MitmRejectImg
	MitmRejectDict
	MitmRejectArray

	Mitm302
	Mitm307

	MitmRequestHeader
	MitmRequestBody

	MitmResponseHeader
	MitmResponseBody
)

type RewriteType int

// UnmarshalYAML unserialize RewriteType with yaml
func (e *RewriteType) UnmarshalYAML(unmarshal func(any) error) error {
	var tp string
	if err := unmarshal(&tp); err != nil {
		return err
	}
	mode, exist := RewriteTypeMapping[tp]
	if !exist {
		return errors.New("invalid MITM Action")
	}
	*e = mode
	return nil
}

// MarshalYAML serialize RewriteType with yaml
func (e RewriteType) MarshalYAML() (any, error) {
	return e.String(), nil
}

// UnmarshalJSON unserialize RewriteType with json
func (e *RewriteType) UnmarshalJSON(data []byte) error {
	var tp string
	json.Unmarshal(data, &tp)
	mode, exist := RewriteTypeMapping[tp]
	if !exist {
		return errors.New("invalid MITM Action")
	}
	*e = mode
	return nil
}

// MarshalJSON serialize RewriteType with json
func (e RewriteType) MarshalJSON() ([]byte, error) {
	return json.Marshal(e.String())
}

func (rt RewriteType) String() string {
	switch rt {
	case MitmReject:
		return "reject" // 404
	case MitmReject200:
		return "reject-200"
	case MitmRejectImg:
		return "reject-img"
	case MitmRejectDict:
		return "reject-dict"
	case MitmRejectArray:
		return "reject-array"
	case Mitm302:
		return "302"
	case Mitm307:
		return "307"
	case MitmRequestHeader:
		return "request-header"
	case MitmRequestBody:
		return "request-body"
	case MitmResponseHeader:
		return "response-header"
	case MitmResponseBody:
		return "response-body"
	default:
		return "Unknown"
	}
}

type Rewrite interface {
	ID() string
	URLRegx() *regexp.Regexp
	RuleType() RewriteType
	RuleRegx() *regexp.Regexp
	RulePayload() string
	ReplaceURLPayload([]string) string
	ReplaceSubPayload(string) string
}

type RewriteRule interface {
	SearchInRequest(func(Rewrite) bool) bool
	SearchInResponse(func(Rewrite) bool) bool
}