Add apikey

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:

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
 	}

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 {

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 {

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 {