diff --git a/docker-compose.yaml b/docker-compose.yaml index a7d49c0..25b0d7c 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -3,19 +3,15 @@ version: "3" services: wg: build: . - #image: ngoduykhanh/wireguard-ui:latest + image: wg-ui:latest container_name: wgui cap_add: - NET_ADMIN network_mode: host environment: - - SENDGRID_API_KEY - - EMAIL_FROM_ADDRESS - - EMAIL_FROM_NAME - - SESSION_SECRET + - SESSION_SECRET=vnsjdvb9134f39hvn9249 - WGUI_USERNAME=alpha - - WGUI_PASSWORD=this-unusual-password - - WG_CONF_TEMPLATE + - WGUI_PASSWORD=Negjqgfhjkm1! - WGUI_MANAGE_START=false - WGUI_MANAGE_RESTART=false logging: diff --git a/handler/session.go b/handler/session.go index 9975e0d..fb2c634 100644 --- a/handler/session.go +++ b/handler/session.go @@ -23,12 +23,32 @@ func ValidSession(next echo.HandlerFunc) echo.HandlerFunc { } } +func ProtectedHandler(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + sess, err := session.Get("session", c) + if err != nil { + return err + } + + apiKey, ok := sess.Values["api_key"].(string) + if !ok || apiKey != util.ApiKey { + return echo.NewHTTPError(http.StatusUnauthorized, "Invalid API key") + } + // Handle the request for authenticated users + return next(c) + } +} + func isValidSession(c echo.Context) bool { if util.DisableLogin { return true } sess, _ := session.Get("session", c) cookie, err := c.Cookie("session_token") + apiKey, ok := sess.Values["api_key"].(string) + if ok && apiKey != util.ApiKey { + return false + } if err != nil || sess.Values["session_token"] != cookie.Value { return false } diff --git a/main.go b/main.go index cbfa8b7..7c8888e 100644 --- a/main.go +++ b/main.go @@ -40,6 +40,7 @@ var ( flagSessionSecret string flagWgConfTemplate string flagBasePath string + flagApiKey string ) const ( @@ -69,6 +70,7 @@ func init() { flag.StringVar(&flagSessionSecret, "session-secret", util.LookupEnvOrString("SESSION_SECRET", flagSessionSecret), "The key used to encrypt session cookies.") flag.StringVar(&flagWgConfTemplate, "wg-conf-template", util.LookupEnvOrString("WG_CONF_TEMPLATE", flagWgConfTemplate), "Path to custom wg.conf template.") flag.StringVar(&flagBasePath, "base-path", util.LookupEnvOrString("BASE_PATH", flagBasePath), "The base path of the URL") + flag.StringVar(&flagApiKey, "api-key", util.LookupEnvOrString("WGUI_API_KEY", ""), "Specify API key for auth") flag.Parse() // update runtime config @@ -87,7 +89,7 @@ func init() { util.SessionSecret = []byte(flagSessionSecret) util.WgConfTemplate = flagWgConfTemplate util.BasePath = util.ParseBasePath(flagBasePath) - + util.ApiKey = flagApiKey // print app information fmt.Println("Wireguard UI") fmt.Println("App Version\t:", appVersion) @@ -105,6 +107,7 @@ func init() { fmt.Println("Base path\t:", util.BasePath+"/") } + func main() { db, err := jsondb.New("./db") if err != nil { diff --git a/router/router.go b/router/router.go index 9aeaf1b..1446b69 100644 --- a/router/router.go +++ b/router/router.go @@ -47,11 +47,31 @@ func (t *TemplateRegistry) Render(w io.Writer, name string, data interface{}, c return tmpl.ExecuteTemplate(w, "base.html", data) } +func apiKeyMiddleware(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + apiKey := c.Request().Header.Get("X-API-Key") + if apiKey == "" { + apiKey = c.QueryParam("api_key") + } + sess, err := session.Get("session", c) + if err != nil { + return err + } + sess.Values["api_key"] = apiKey + err = sess.Save(c.Request(), c.Response()) + if err != nil { + return err + } + return next(c) + } +} // New function func New(tmplBox *rice.Box, extraData map[string]string, secret []byte) *echo.Echo { e := echo.New() - e.Use(session.Middleware(sessions.NewCookieStore(secret))) + store := sessions.NewCookieStore(secret) + e.Use(session.Middleware(store)) + e.Use(apiKeyMiddleware) // read html template file to string tmplBaseString, err := tmplBox.String("base.html") if err != nil { diff --git a/util/config.go b/util/config.go index 7f5d221..447d913 100644 --- a/util/config.go +++ b/util/config.go @@ -19,6 +19,7 @@ var ( SessionSecret []byte WgConfTemplate string BasePath string + ApiKey string ) const ( @@ -49,6 +50,7 @@ const ( DefaultClientExtraAllowedIpsEnvVar = "WGUI_DEFAULT_CLIENT_EXTRA_ALLOWED_IPS" DefaultClientUseServerDNSEnvVar = "WGUI_DEFAULT_CLIENT_USE_SERVER_DNS" DefaultClientEnableAfterCreationEnvVar = "WGUI_DEFAULT_CLIENT_ENABLE_AFTER_CREATION" + DefaultApiKeyEnvVar = "WGUI_API_KEY" ) func ParseBasePath(basePath string) string {