diff --git a/src/internal/config/config.go b/src/internal/config/config.go index a2da344..974dead 100644 --- a/src/internal/config/config.go +++ b/src/internal/config/config.go @@ -1,10 +1,20 @@ package config -import "flag" +import ( + "flag" + "fmt" +) + +const ( + ServerModeSocks5 = "socks5" + ServerModeTProxy = "tproxy" +) type Config struct { + ServerMode string BindAddr string Port int + ListenAddr string LogLevel string PayloadUA string UAPattern string @@ -13,15 +23,17 @@ type Config struct { func Parse() (*Config, bool) { var ( - bindAddr string - port int - loglevel string - payloadUA string - uaPattern string - partial bool - showVer bool + serverMode string + bindAddr string + port int + loglevel string + payloadUA string + uaPattern string + partial bool + showVer bool ) + flag.StringVar(&serverMode, "m", ServerModeSocks5, "server mode: socks5 or tproxy (default: socks5)") flag.StringVar(&bindAddr, "b", "127.0.0.1", "bind address (default: 127.0.0.1)") flag.IntVar(&port, "p", 1080, "port") flag.StringVar(&payloadUA, "f", "FFF", "User-Agent") @@ -32,8 +44,10 @@ func Parse() (*Config, bool) { flag.Parse() cfg := &Config{ + ServerMode: serverMode, BindAddr: bindAddr, Port: port, + ListenAddr: fmt.Sprintf("%s:%d", bindAddr, port), LogLevel: loglevel, PayloadUA: payloadUA, UAPattern: uaPattern, diff --git a/src/internal/log/log.go b/src/internal/log/log.go index 8cfcffc..b7b0d77 100644 --- a/src/internal/log/log.go +++ b/src/internal/log/log.go @@ -63,9 +63,9 @@ func SetLogConf(level string) { } } -func LogHeader(version string, addr string, cfg *config.Config) { +func LogHeader(version string, cfg *config.Config) { logrus.Info("UA3F v" + version) - logrus.Infof("Listen on %s", addr) + logrus.Infof("Listen on %s", cfg.ListenAddr) logrus.Infof("User-Agent: %s", cfg.PayloadUA) logrus.Infof("User-Agent Regex Pattern: '%s'", cfg.UAPattern) logrus.Infof("Enable Partial Replace: %v", cfg.EnablePartialReplace) diff --git a/src/internal/server/server.go b/src/internal/server/server.go new file mode 100644 index 0000000..bd6f73e --- /dev/null +++ b/src/internal/server/server.go @@ -0,0 +1,22 @@ +package server + +import ( + "fmt" + + "github.com/sunbk201/ua3f/internal/config" + "github.com/sunbk201/ua3f/internal/rewrite" + "github.com/sunbk201/ua3f/internal/server/socks5" +) + +type Server interface { + Start() error +} + +func NewServer(cfg *config.Config, rw *rewrite.Rewriter) (Server, error) { + switch cfg.ServerMode { + case config.ServerModeSocks5: + return socks5.New(cfg, rw), nil + default: + return nil, fmt.Errorf("unknown server mode: %s", cfg.ServerMode) + } +} diff --git a/src/internal/server/socks5/socks5.go b/src/internal/server/socks5/socks5.go index 7ffdbe0..5f9ca9c 100644 --- a/src/internal/server/socks5/socks5.go +++ b/src/internal/server/socks5/socks5.go @@ -34,24 +34,22 @@ var ( // Server is a minimal SOCKS5 server that delegates HTTP UA rewriting to Rewriter. type Server struct { - cfg *config.Config - rw *rewrite.Rewriter - listener net.Listener - ListenAddr string + cfg *config.Config + rw *rewrite.Rewriter + listener net.Listener } // New returns a new Server with given config, rewriter, and version string. func New(cfg *config.Config, rw *rewrite.Rewriter) *Server { return &Server{ - cfg: cfg, - rw: rw, - ListenAddr: fmt.Sprintf("%s:%d", cfg.BindAddr, cfg.Port), + cfg: cfg, + rw: rw, } } // Start begins listening for SOCKS5 clients. func (s *Server) Start() (err error) { - if s.listener, err = net.Listen("tcp", s.ListenAddr); err != nil { + if s.listener, err = net.Listen("tcp", s.cfg.ListenAddr); err != nil { return fmt.Errorf("listen failed: %w", err) } diff --git a/src/main.go b/src/main.go index 092ebb5..ce1c2ba 100644 --- a/src/main.go +++ b/src/main.go @@ -8,7 +8,7 @@ import ( "github.com/sunbk201/ua3f/internal/config" "github.com/sunbk201/ua3f/internal/log" "github.com/sunbk201/ua3f/internal/rewrite" - "github.com/sunbk201/ua3f/internal/server/socks5" + "github.com/sunbk201/ua3f/internal/server" ) const version = "0.9.0" @@ -26,8 +26,13 @@ func main() { logrus.Fatal(err) } - srv := socks5.New(cfg, rw) - log.LogHeader(version, srv.ListenAddr, cfg) + srv, err := server.NewServer(cfg, rw) + if err != nil { + logrus.Fatal(err) + } + + log.LogHeader(version, cfg) + if err := srv.Start(); err != nil { logrus.Fatal(err) }