update cli flags to load ruleset

This commit is contained in:
Kevin Pham
2023-12-05 21:28:53 -06:00
parent 5a7fe8a70a
commit 6192373587
4 changed files with 28 additions and 18 deletions

View File

@@ -12,7 +12,7 @@ tmp_dir = "tmp"
exclude_regex = ["_test.go"] exclude_regex = ["_test.go"]
exclude_unchanged = false exclude_unchanged = false
follow_symlink = false follow_symlink = false
full_bin = "./tmp/main --ruleset ./ruleset.yaml" full_bin = "./tmp/main --ruleset ./ruleset_v2.yaml"
include_dir = [] include_dir = []
include_ext = ["go", "tpl", "tmpl", "yaml", "html", "js"] include_ext = ["go", "tpl", "tmpl", "yaml", "html", "js"]
include_file = [] include_file = []

View File

@@ -10,6 +10,7 @@ import (
"ladder/handlers" "ladder/handlers"
"ladder/internal/cli" "ladder/internal/cli"
"ladder/proxychain/requestmodifiers/bot" "ladder/proxychain/requestmodifiers/bot"
"ladder/proxychain/ruleset"
"github.com/akamensky/argparse" "github.com/akamensky/argparse"
"github.com/gofiber/fiber/v2" "github.com/gofiber/fiber/v2"
@@ -116,6 +117,18 @@ func main() {
*prefork = true *prefork = true
} }
var rs ruleset_v2.IRuleset
switch {
case *ruleset != "":
rs, err = ruleset_v2.NewRuleset(*ruleset)
if err != nil {
fmt.Printf("ERROR: failed to load ruleset from %s\n", *ruleset)
}
case os.Getenv("RULESET") != "":
rs = ruleset_v2.NewRulesetFromEnv()
}
engine := html.New("./handlers", ".html") engine := html.New("./handlers", ".html")
engine.AddFunc( engine.AddFunc(
// add unescape function // add unescape function
@@ -145,6 +158,11 @@ func main() {
}) })
} }
proxyOpts := &handlers.ProxyOptions{
Verbose: *verbose,
Ruleset: rs,
}
app.Get("/", handlers.Form) app.Get("/", handlers.Form)
app.Get("styles.css", handlers.Styles) app.Get("styles.css", handlers.Styles)
@@ -152,11 +170,6 @@ func main() {
app.Get("ruleset", handlers.Ruleset) app.Get("ruleset", handlers.Ruleset)
app.Get("raw/*", handlers.Raw) app.Get("raw/*", handlers.Raw)
proxyOpts := &handlers.ProxyOptions{
Verbose: *verbose,
RulesetPath: *ruleset,
}
app.Get("api/content/*", handlers.NewAPIContentHandler("api/outline/*", proxyOpts)) app.Get("api/content/*", handlers.NewAPIContentHandler("api/outline/*", proxyOpts))
app.Get("outline/*", handlers.NewOutlineHandler("outline/*", proxyOpts)) app.Get("outline/*", handlers.NewOutlineHandler("outline/*", proxyOpts))

View File

@@ -10,8 +10,8 @@ import (
) )
type ProxyOptions struct { type ProxyOptions struct {
RulesetPath string Ruleset ruleset_v2.IRuleset
Verbose bool Verbose bool
} }
func NewProxySiteHandler(opts *ProxyOptions) fiber.Handler { func NewProxySiteHandler(opts *ProxyOptions) fiber.Handler {
@@ -25,10 +25,6 @@ func NewProxySiteHandler(opts *ProxyOptions) fiber.Handler {
rs = r rs = r
} }
*/ */
rs, err := ruleset_v2.NewRuleset("ruleset_v2.yaml")
if err != nil {
panic(err)
}
return func(c *fiber.Ctx) error { return func(c *fiber.Ctx) error {
proxychain := proxychain. proxychain := proxychain.
@@ -59,7 +55,7 @@ func NewProxySiteHandler(opts *ProxyOptions) fiber.Handler {
) )
// load ruleset // load ruleset
rule, exists := rs.GetRule(proxychain.Request.URL) rule, exists := opts.Ruleset.GetRule(proxychain.Request.URL)
if exists { if exists {
proxychain.AddOnceRequestModifications(rule.RequestModifications...) proxychain.AddOnceRequestModifications(rule.RequestModifications...)
proxychain.AddOnceResponseModifications(rule.ResponseModifications...) proxychain.AddOnceResponseModifications(rule.ResponseModifications...)

View File

@@ -18,9 +18,10 @@ import (
) )
type IRuleset interface { type IRuleset interface {
HasRule(url url.URL) bool
GetRule(url url.URL) (rule Rule, exists bool)
YAML() (string, error) YAML() (string, error)
JSON() (string, error)
HasRule(url *url.URL) bool
GetRule(url *url.URL) (rule *Rule, exists bool)
} }
type Ruleset struct { type Ruleset struct {
@@ -318,7 +319,7 @@ func (rs *Ruleset) loadRulesFromRemoteFile(rulesURL string) error {
// ================= utility methods ========================== // ================= utility methods ==========================
// YAML returns the ruleset as a Yaml string // YAML returns the ruleset as a Yaml string
func (rs *Ruleset) YAML() (string, error) { func (rs Ruleset) YAML() (string, error) {
y, err := yaml.Marshal(rs) y, err := yaml.Marshal(rs)
if err != nil { if err != nil {
return "", err return "", err
@@ -333,8 +334,8 @@ func (rs *Ruleset) YAML() (string, error) {
return x, nil return x, nil
} }
// YAML returns the ruleset as a JSON string // JSON returns the ruleset as a JSON string
func (rs *Ruleset) JSON() (string, error) { func (rs Ruleset) JSON() (string, error) {
j, err := json.Marshal(rs) j, err := json.Marshal(rs)
if err != nil { if err != nil {
return "", err return "", err