waf.go 32.1 KB
Newer Older
qiuqunfeng's avatar
qiuqunfeng committed
1 2 3
package service

import (
qiuqunfeng's avatar
qiuqunfeng committed
4
	"bytes"
qiuqunfeng's avatar
qiuqunfeng committed
5
	"context"
qiuqunfeng's avatar
commit  
qiuqunfeng committed
6
	"crypto/tls"
qiuqunfeng's avatar
qiuqunfeng committed
7
	"encoding/json"
qiuqunfeng's avatar
commit  
qiuqunfeng committed
8
	"fmt"
9
	"io"
qiuqunfeng's avatar
qiuqunfeng committed
10
	"net/http"
11
	"net/url"
qiuqunfeng's avatar
commit  
qiuqunfeng committed
12
	"os"
qiuqunfeng's avatar
qiuqunfeng committed
13 14
	"strconv"
	"strings"
15
	"sync"
qiuqunfeng's avatar
qiuqunfeng committed
16

17 18
	jsoniter "github.com/json-iterator/go"
	"github.com/olivere/elastic/v7"
qiuqunfeng's avatar
commit  
qiuqunfeng committed
19
	"github.com/rs/zerolog/log"
qiuqunfeng's avatar
commit  
qiuqunfeng committed
20
	"gitlab.com/tensorsecurity-rd/waf-console/internal/model"
qiuqunfeng's avatar
qiuqunfeng committed
21
	"gitlab.com/tensorsecurity-rd/waf-console/internal/utils"
qiuqunfeng's avatar
qiuqunfeng committed
22 23
	"gitlab.com/tensorsecurity-rd/waf-console/pkg/apis/waf.security.io/v1alpha1"
	"gorm.io/gorm"
24
	corev1 "k8s.io/api/core/v1"
qiuqunfeng's avatar
qiuqunfeng committed
25
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
26
	"k8s.io/apimachinery/pkg/util/sets"
qiuqunfeng's avatar
qiuqunfeng committed
27 28 29
)

type wafService struct {
qiuqunfeng's avatar
qiuqunfeng committed
30 31
	clusterClientManager *utils.ClusterClientManager
	db                   *gorm.DB
32
	gatewayUrl           string
33
	elasticClient        *elastic.Client
qiuqunfeng's avatar
qiuqunfeng committed
34 35
}

36 37
func NewWafService(clusterClientManager *utils.ClusterClientManager, db *gorm.DB, gatewayUrl string, elasticClient *elastic.Client) Service {
	return &wafService{clusterClientManager: clusterClientManager, db: db, gatewayUrl: gatewayUrl, elasticClient: elasticClient}
qiuqunfeng's avatar
qiuqunfeng committed
38 39
}

qiuqunfeng's avatar
qiuqunfeng committed
40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57
func (s *wafService) GetWaf(ctx context.Context, regionCode, namespace, gatewayName string) (*WafService, error) {
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ? AND region_code = ? AND namespace = ?", gatewayName, regionCode, namespace).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			// Create new WAF service record if not found
			wafService = &model.WafService{
				RegionCode:  regionCode,
				Namespace:   namespace,
				GatewayName: gatewayName,
				Mode:        string(WafModeAlert),
			}
			if err := s.db.Create(wafService).Error; err != nil {
				return nil, fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return nil, fmt.Errorf("failed to query WAF service: %v", err)
		}
qiuqunfeng's avatar
qiuqunfeng committed
58
	}
59 60 61 62 63 64 65 66 67 68 69 70 71 72
	listenerWafs, err := s.ListListenerWafStatus(ctx, &GatewateInfo{
		GatewayName: gatewayName,
		Namespace:   namespace,
		RegionCode:  regionCode,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to list listener WAF status: %v", err)
	}
	listeners := []string{}
	for _, listener := range listenerWafs {
		hosts := strings.Join(listener.Hosts, "@")
		listeners = append(listeners, fmt.Sprintf("%s-%d", hosts, listener.Port))
	}

qiuqunfeng's avatar
qiuqunfeng committed
73 74 75 76 77
	return &WafService{
		GatewayName: wafService.GatewayName,
		Mode:        wafService.Mode,
		RuleNum:     wafService.RuleNum,
		AttackNum:   wafService.AttackNum,
78
		Listeners:   listeners,
qiuqunfeng's avatar
qiuqunfeng committed
79
	}, nil
qiuqunfeng's avatar
qiuqunfeng committed
80 81
}

82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130
func (s *wafService) GetWafGatewayInfo(ctx context.Context, req *GetWafGatewayInfoReq) (*WafService, error) {
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			httpRequst := http.Request{
				Method: http.MethodPost,
				URL:    &url.URL{Scheme: "https", Host: "console.tensorsecurity.com", Path: "/api/v1/waf/gateway"},
				Header: http.Header{
					"Cookie": []string{req.Cookie},
				},
				Body: io.NopCloser(strings.NewReader(fmt.Sprintf(`{"gateway_name": "%s", "namespace": "%s", "region_code": "%s"}`, req.GatewayName, req.Namespace, req.RegionCode))),
			}
			resp, err := http.DefaultClient.Do(&httpRequst)
			if err != nil {
				return nil, fmt.Errorf("failed to get WAF service: %v", err)
			}
			defer resp.Body.Close()
			body, err := io.ReadAll(resp.Body)
			if err != nil {
				return nil, fmt.Errorf("failed to read WAF service: %v", err)
			}
			var wafService model.WafService
			err = json.Unmarshal(body, &wafService)
			if err != nil {
				return nil, fmt.Errorf("failed to unmarshal WAF service: %v", err)
			}
			wafService.ID = 0
			wafService.RuleCategoryStatus = nil
			wafService.RuleNum = 0
			wafService.AttackNum = 0
			// wafService.Host = model.HostList([]string{"*"})
			wafService.Mode = string(WafModeAlert)
			err = s.db.Create(wafService).Error
			if err != nil {
				return nil, fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return nil, fmt.Errorf("failed to query WAF service: %v", err)
		}
	}
	return &WafService{
		GatewayName: wafService.GatewayName,
		Mode:        wafService.Mode,
		RuleNum:     wafService.RuleNum,
		AttackNum:   wafService.AttackNum,
	}, nil
}

qiuqunfeng's avatar
commit  
qiuqunfeng committed
131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192
func (s *wafService) getRulesForService(req *CreateWafReq) ([]v1alpha1.Rule, error) {
	rules := []v1alpha1.Rule{}
	ruleCategories := []model.WafRuleCategory{}
	if err := s.db.Model(&model.WafRuleCategory{}).Where("status = ?", 0).Find(&ruleCategories).Error; err != nil {
		return nil, fmt.Errorf("failed to get rule categories: %v", err)
	}

	// Get existing WAF service config if any
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			// Create new WAF service record if not found
			wafService = &model.WafService{
				RegionCode:  req.RegionCode,
				Namespace:   req.Namespace,
				GatewayName: req.GatewayName,
				Mode:        string(WafModeAlert),
			}
			if err := s.db.Create(wafService).Error; err != nil {
				return nil, fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return nil, fmt.Errorf("failed to query WAF service: %v", err)
		}
	}

	// Determine which rule categories to enable
	var enabledCategories []model.WafRuleCategory

	if wafService.RuleCategoryStatus != nil && len(wafService.RuleCategoryStatus.CategoryID) > 0 {
		// Only include categories not already enabled
		for _, category := range ruleCategories {
			for _, id := range wafService.RuleCategoryStatus.CategoryID {
				if id == category.CategoryID {
					enabledCategories = append(enabledCategories, category)
					continue
				}
			}
		}
	} else {
		// Enable all categories if none specified
		enabledCategories = ruleCategories
	}

	for _, category := range enabledCategories {
		for _, rule := range category.Rules {
			rules = append(rules, v1alpha1.Rule{
				ID:          rule.ID,
				Level:       rule.Level,
				Name:        rule.Name,
				Type:        rule.Type,
				Description: rule.Description,
				Expr:        rule.Expr,
				Mode:        rule.Mode,
			})
		}
	}

	return rules, nil
}

qiuqunfeng's avatar
qiuqunfeng committed
193 194 195
func (s *wafService) CreateWaf(ctx context.Context, req *CreateWafReq) (*WafService, error) {
	// Create the WAF service resource
	name := fmt.Sprintf("%s-%d", req.GatewayName, req.Port)
qiuqunfeng's avatar
qiuqunfeng committed
196 197
	service := &v1alpha1.Service{
		ObjectMeta: metav1.ObjectMeta{
qiuqunfeng's avatar
qiuqunfeng committed
198
			Name:      name,
qiuqunfeng's avatar
qiuqunfeng committed
199
			Namespace: req.Namespace,
qiuqunfeng's avatar
qiuqunfeng committed
200 201 202
			Labels: map[string]string{
				"apigateway_name": req.GatewayName,
			},
qiuqunfeng's avatar
qiuqunfeng committed
203 204 205
		},
		Spec: v1alpha1.ServiceSpec{
			HostNames:   req.Host,
206
			ServiceName: req.ListenerName,
qiuqunfeng's avatar
qiuqunfeng committed
207 208
			Port:        req.Port,
			Workload: v1alpha1.WorkloadRef{
209 210 211 212
				Kind:       req.GatewayName,
				Name:       req.ListenerName,
				Namespace:  req.Namespace,
				ClusterKey: req.RegionCode,
qiuqunfeng's avatar
qiuqunfeng committed
213
			},
qiuqunfeng's avatar
qiuqunfeng committed
214 215 216 217 218 219 220
			Uri: &v1alpha1.StringMatch{
				Prefix: "/",
			},
			LogConfig: &v1alpha1.LogConfig{
				Enable: 1,
				Level:  "info",
			},
221 222
			Mode:      string(req.Mode),
			ServiceID: req.ServiceID,
qiuqunfeng's avatar
qiuqunfeng committed
223 224
		},
	}
qiuqunfeng's avatar
qiuqunfeng committed
225

226
	rules, err := s.getRulesForService(req)
qiuqunfeng's avatar
qiuqunfeng committed
227
	if err != nil {
228
		return nil, fmt.Errorf("failed to get rules for service: %v", err)
qiuqunfeng's avatar
qiuqunfeng committed
229
	}
230
	service.Spec.Rules = rules
qiuqunfeng's avatar
qiuqunfeng committed
231

232 233 234 235
	if len(service.Spec.Rules) == 0 {
		return nil, fmt.Errorf("cannot create WAF service with no rules")
	}

qiuqunfeng's avatar
qiuqunfeng committed
236 237 238
	// Create the WAF service in Kubernetes
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
239
		return nil, fmt.Errorf("failed to get cluster client for region %s", req.RegionCode)
qiuqunfeng's avatar
qiuqunfeng committed
240
	}
241
	if _, err := client.Versioned.WafV1alpha1().Services(req.Namespace).Create(ctx, service, metav1.CreateOptions{}); err != nil {
qiuqunfeng's avatar
qiuqunfeng committed
242
		return nil, fmt.Errorf("failed to create WAF service: %v", err)
qiuqunfeng's avatar
qiuqunfeng committed
243 244
	}

245 246 247 248 249 250
	return &WafService{
		GatewayName: req.GatewayName,
		Mode:        service.Spec.Mode,
		RuleNum:     len(service.Spec.Rules),
		AttackNum:   0,
	}, nil
qiuqunfeng's avatar
qiuqunfeng committed
251
}
qiuqunfeng's avatar
commit  
qiuqunfeng committed
252

qiuqunfeng's avatar
qiuqunfeng committed
253 254 255
func (s *wafService) DeleteListenerWaf(ctx context.Context, req *DeleteListenerReq) error {
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
256
		return fmt.Errorf("failed to get cluster client for region %s", req.RegionCode)
qiuqunfeng's avatar
qiuqunfeng committed
257 258
	}
	name := fmt.Sprintf("%s-%d", req.GatewayName, req.Port)
259
	if err := client.Versioned.WafV1alpha1().Services(req.Namespace).Delete(ctx, name, metav1.DeleteOptions{}); err != nil {
qiuqunfeng's avatar
qiuqunfeng committed
260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289
		return fmt.Errorf("failed to delete WAF service: %v", err)
	}
	return nil
}

func (s *wafService) UpdateMode(ctx context.Context, req *UpdateModeReq) (*WafService, error) {
	// Check if WAF service exists
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ?", req.GatewayName).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			// Create new WAF service record if not found
			wafService = &model.WafService{
				RegionCode:  req.RegionCode,
				Namespace:   req.Namespace,
				GatewayName: req.GatewayName,
				Mode:        string(req.Mode),
			}
			if err := s.db.Create(wafService).Error; err != nil {
				return nil, fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return nil, fmt.Errorf("failed to query WAF service: %v", err)
		}
	} else {
		// Update mode if service exists
		if err := s.db.Model(wafService).Update("mode", string(req.Mode)).Error; err != nil {
			return nil, fmt.Errorf("failed to update WAF service mode: %v", err)
		}
	}
290 291 292 293 294
	// Update mode for each listener
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		return nil, fmt.Errorf("failed to get cluster client for region %s", req.RegionCode)
	}
295
	listenerList, err := client.Versioned.WafV1alpha1().Services(req.Namespace).List(ctx, metav1.ListOptions{LabelSelector: fmt.Sprintf("apigateway_name=%s", req.GatewayName)})
296 297 298
	if err != nil {
		return nil, fmt.Errorf("failed to get listener list: %v", err)
	}
299
	var wg sync.WaitGroup
300
	for _, listener := range listenerList.Items {
301
		wg.Add(1)
302
		listener := listener // Create new variable for goroutine
303
		listener.Spec.Mode = string(req.Mode)
304
		go func() {
305
			defer wg.Done()
306
			log.Info().Msgf("update WAF service mode: %v", listener.Name)
307
			_, err := client.Versioned.WafV1alpha1().Services(req.Namespace).Update(ctx, &listener, metav1.UpdateOptions{})
308 309 310 311 312
			if err != nil {
				log.Error().Msgf("failed to update WAF service mode: %v", err)
			}
		}()
	}
313
	wg.Wait()
314 315 316 317
	return &WafService{
		GatewayName: req.GatewayName,
		Mode:        string(req.Mode),
	}, nil
qiuqunfeng's avatar
commit  
qiuqunfeng committed
318
}
qiuqunfeng's avatar
commit  
qiuqunfeng committed
319

320 321 322 323 324 325 326 327
// func (s *wafService) GetRuleCategories(ctx context.Context) ([]WafRuleCategory, error) {
// 	var categories []WafRuleCategory
// 	err := s.db.Table("waf_rule_categories").Find(&categories).Error
// 	if err != nil {
// 		return nil, err
// 	}
// 	return categories, nil
// }
qiuqunfeng's avatar
commit  
qiuqunfeng committed
328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348

func (s *wafService) GetRules(ctx context.Context, categoryID string) ([]WafRule, error) {
	var rules []WafRule
	err := s.db.Table("waf_rules").Where("category_id = ?", categoryID).Find(&rules).Error
	if err != nil {
		return nil, err
	}
	return rules, nil
}

func (s *wafService) GetRule(ctx context.Context, ruleID int) (*WafRule, error) {
	var rule WafRule
	err := s.db.Table("waf_rules").Where("id = ?", ruleID).First(&rule).Error
	if err != nil {
		return nil, err
	}
	return &rule, nil
}

func (s *wafService) SaveRuleCategoryToDB(ctx context.Context) error {
	var categories []WafRuleCategory
349
	jsonFile, err := os.ReadFile("rules/waf-rules.json")
qiuqunfeng's avatar
commit  
qiuqunfeng committed
350 351 352 353
	if err != nil {
		return fmt.Errorf("error reading yaml file: %v", err)
	}

354 355
	// err = yaml.Unmarshal(yamlFile, &categories)
	err = json.Unmarshal(jsonFile, &categories)
qiuqunfeng's avatar
commit  
qiuqunfeng committed
356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376
	if err != nil {
		return fmt.Errorf("error unmarshaling yaml: %v", err)
	}

	for _, category := range categories {
		rules := []model.WafRule{}
		for _, rule := range category.Rules {
			rules = append(rules, model.WafRule{
				ID:          rule.ID,
				CategoryID:  category.CategoryID,
				Level:       rule.Level,
				Name:        rule.Name,
				Type:        rule.Type,
				Description: rule.Description,
				Expr:        rule.Expr,
				Mode:        rule.Mode,
			})
		}
		model := model.WafRuleCategory{
			CategoryID:    category.CategoryID,
			Status:        category.Status,
377 378
			CategoryEN:    category.Category.EN,
			CategoryZH:    category.Category.Zh,
qiuqunfeng's avatar
commit  
qiuqunfeng committed
379 380 381 382
			DescriptionEN: category.Description.EN,
			DescriptionZH: category.Description.Zh,
			Rules:         model.RuleList(rules),
		}
qiuqunfeng's avatar
commit  
qiuqunfeng committed
383 384 385 386
		err = s.db.Table("waf_rule_categories").Create(&model).Error
		if err != nil {
			return err
		}
qiuqunfeng's avatar
commit  
qiuqunfeng committed
387 388 389 390
	}

	return nil
}
qiuqunfeng's avatar
qiuqunfeng committed
391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406

func (s *wafService) DeleteListener(ctx context.Context, req *DeleteListenerReq) error {
	listener := &model.GatewayListener{}
	err := s.db.Model(&model.GatewayListener{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).First(listener).Error
	if err != nil {
		return err
	}

	err = s.db.Model(&model.GatewayListener{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", req.GatewayName, req.Namespace, req.RegionCode).Delete(listener).Error
	if err != nil {
		return err
	}

	return nil
}

407 408 409 410 411 412
func (s *wafService) getServiceID(ctx context.Context, gatewayName, namespace, regionCode string) (uint32, error) {
	service := &model.WafService{}
	err := s.db.WithContext(ctx).Model(&model.WafService{}).Where("gateway_name = ? AND namespace = ? AND region_code = ?", gatewayName, namespace, regionCode).First(service).Error
	if err != nil {
		return 0, err
	}
qiuqunfeng's avatar
qiuqunfeng committed
413

414 415 416 417
	return uint32(service.ID), nil
}

func (s *wafService) EnableListenerWaf(ctx context.Context, req *EnableListenerWafReq) error {
418 419
	if req.Enable {
		log.Info().Msgf("Create WAF for listener %s", req.GatewayName)
420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438
		serviceID, err := s.getServiceID(ctx, req.GatewayName, req.Namespace, req.RegionCode)
		if err != nil {
			if err == gorm.ErrRecordNotFound {
				service := &model.WafService{
					GatewayName: req.GatewayName,
					Namespace:   req.Namespace,
					RegionCode:  req.RegionCode,
					Mode:        string(req.Mode),
				}
				if err := s.db.Model(&model.WafService{}).Create(&service).Error; err != nil {
					return err
				}
				serviceID = uint32(service.ID)
			} else {
				return err
			}
		}

		_, err = s.CreateWaf(ctx, &CreateWafReq{
qiuqunfeng's avatar
qiuqunfeng committed
439 440 441 442 443
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
444 445 446 447
			Port:         uint32(req.Port),
			Host:         req.Hosts,
			Mode:         req.Mode,
			ListenerName: req.ListenerName,
448
			ServiceID:    serviceID,
qiuqunfeng's avatar
qiuqunfeng committed
449
		})
qiuqunfeng's avatar
commit  
qiuqunfeng committed
450 451 452
		if err != nil {
			return err
		}
qiuqunfeng's avatar
qiuqunfeng committed
453
	} else {
454 455
		log.Info().Msgf("Delete WAF for listener %s", req.GatewayName)
		err := s.DeleteListenerWaf(ctx, &DeleteListenerReq{
qiuqunfeng's avatar
qiuqunfeng committed
456 457 458 459 460 461 462
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
			Port: req.Port,
		})
qiuqunfeng's avatar
commit  
qiuqunfeng committed
463 464 465
		if err != nil {
			return err
		}
qiuqunfeng's avatar
qiuqunfeng committed
466 467 468 469
	}
	return nil
}

470 471 472 473 474 475
func getGatewayNameFromCrn(crn string) string {
	// crn:ucs::apigateway:lf-tst7:214613666997:instance/testaaa
	parts := strings.Split(crn, "/")
	return parts[len(parts)-1]
}

476
func (s *wafService) listListenerFromApiGateway(ctx context.Context, apiGatewayCrn string, regionCode string, cookie string) ([]GatewayRespListenerData, error) {
477 478 479 480 481 482 483
	body, err := json.Marshal(map[string]string{
		"apigateway_crn": apiGatewayCrn,
		"region_code":    regionCode,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request body: %v", err)
	}
484
	request, err := http.NewRequestWithContext(ctx, "POST", "https://csm.console.test.tg.unicom.local/apigatewaymng/listener/lf-tst7/list_listeners", bytes.NewBuffer(body))
485 486 487 488
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %v", err)
	}
	request.Header.Set("Cookie", cookie)
qiuqunfeng's avatar
commit  
qiuqunfeng committed
489 490 491 492 493 494 495 496
	// Create custom transport with TLS config
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true, // Skip certificate verification for test environment
		},
	}
	client := &http.Client{Transport: tr}
	resp, err := client.Do(request)
497 498 499 500 501
	if err != nil {
		return nil, fmt.Errorf("failed to get listener list: %v", err)
	}
	defer resp.Body.Close()

502
	log.Info().Msgf("resp: %v", resp)
503 504 505 506 507 508
	// Parse response
	var response GatewayListenerResponseList

	if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
		return nil, fmt.Errorf("failed to parse listener list: %v", err)
	}
509
	log.Info().Msgf("response: %v", response)
510 511 512
	return response.Data, nil
}

qiuqunfeng's avatar
qiuqunfeng committed
513 514
func (s *wafService) EnableGatewayWaf(ctx context.Context, req *EnableGatewayWafReq) error {
	if req.Enable {
515
		listeners, err := s.listListenerFromApiGateway(ctx, req.ApiGatewayCrn, req.RegionCode, req.Cookie)
qiuqunfeng's avatar
qiuqunfeng committed
516 517 518
		if err != nil {
			return fmt.Errorf("failed to get listener list: %v", err)
		}
qiuqunfeng's avatar
commit  
qiuqunfeng committed
519
		log.Info().Msgf("listeners: %v", listeners)
qiuqunfeng's avatar
qiuqunfeng committed
520 521
		// Create WAF for each listener
		for _, listener := range listeners {
qiuqunfeng's avatar
commit  
qiuqunfeng committed
522 523
			gatewayName := getGatewayNameFromCrn(listener.ApiGatewayCrn)
			namespace := fmt.Sprintf("%s-%s", listener.CreateAccountName, listener.CreateAccountID)
qiuqunfeng's avatar
qiuqunfeng committed
524 525
			if _, err := s.CreateWaf(ctx, &CreateWafReq{
				GatewateInfo: GatewateInfo{
qiuqunfeng's avatar
commit  
qiuqunfeng committed
526 527
					GatewayName: gatewayName,
					Namespace:   namespace,
qiuqunfeng's avatar
qiuqunfeng committed
528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551
					RegionCode:  req.RegionCode,
				},
				Port: uint32(listener.Port),
				Host: listener.Hosts,
			}); err != nil {
				return fmt.Errorf("failed to create WAF for listener %d: %v", listener.Port, err)
			}
		}
	} else {
		s.DeleteGatewayWaf(ctx, &GatewateInfo{
			GatewayName: req.GatewayName,
			Namespace:   req.Namespace,
			RegionCode:  req.RegionCode,
		})
	}
	return nil
}

func (s *wafService) DeleteGatewayWaf(ctx context.Context, req *GatewateInfo) error {
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		return fmt.Errorf("failed to get cluster client")
	}
	labelSelector := fmt.Sprintf("apigateway_name=%s", req.GatewayName)
552
	if err := client.Versioned.WafV1alpha1().Services(req.Namespace).DeleteCollection(ctx, metav1.DeleteOptions{}, metav1.ListOptions{LabelSelector: labelSelector}); err != nil {
qiuqunfeng's avatar
qiuqunfeng committed
553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597
		return fmt.Errorf("failed to delete WAF service: %v", err)
	}
	return nil
}

func (s *wafService) UpdateRule(ctx context.Context, req *RuleRequest) error {
	wafService := &model.WafService{}
	err := s.db.Model(&model.WafService{}).Where("gateway_name = ?", req.GatewayName).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			// Create new WAF service record if not found
			wafService = &model.WafService{
				RegionCode:  req.RegionCode,
				Namespace:   req.Namespace,
				GatewayName: req.GatewayName,
				Mode:        string(WafModeAlert),
				RuleCategoryStatus: &model.RuleCategoryStatus{
					CategoryID: req.CategoryID,
					Status:     req.Status,
				},
			}
			if err := s.db.Create(wafService).Error; err != nil {
				return fmt.Errorf("failed to create WAF service: %v", err)
			}
		} else {
			return fmt.Errorf("failed to query WAF service: %v", err)
		}
	} else {
		// Update mode if service exists
		if err := s.db.Model(wafService).Update("rule_category_status", model.RuleCategoryStatus{
			CategoryID: req.CategoryID,
			Status:     req.Status,
		}).Error; err != nil {
			return fmt.Errorf("failed to update WAF service mode: %v", err)
		}
	}
	return nil
}

func (s *wafService) ListListenerWafStatus(ctx context.Context, req *GatewateInfo) ([]*GatewayListener, error) {
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		return nil, fmt.Errorf("failed to get cluster client")
	}

598
	listenerList, err := client.Versioned.WafV1alpha1().Services(req.Namespace).List(ctx, metav1.ListOptions{LabelSelector: fmt.Sprintf("apigateway_name=%s", req.GatewayName)})
qiuqunfeng's avatar
qiuqunfeng committed
599 600 601 602 603 604 605 606
	if err != nil {
		return nil, fmt.Errorf("failed to get listener list: %v", err)
	}

	listenerStatusList := []*GatewayListener{}
	for _, listener := range listenerList.Items {
		n := strings.LastIndex(listener.Name, "-")
		if n == -1 {
qiuqunfeng's avatar
commit  
qiuqunfeng committed
607
			return nil, fmt.Errorf("failed to get listener port: %v", listener.Name)
qiuqunfeng's avatar
qiuqunfeng committed
608 609 610 611
		}
		listenerPort := listener.Name[n+1:]
		listenerPortInt, err := strconv.Atoi(listenerPort)
		if err != nil {
qiuqunfeng's avatar
commit  
qiuqunfeng committed
612
			return nil, fmt.Errorf("failed to parse listener port: %v", err)
qiuqunfeng's avatar
qiuqunfeng committed
613 614
		}

615 616
		// hosts := strings.Join(listener.Spec.HostNames, "@")
		// log.Info().Msgf("hosts: %v", hosts)
qiuqunfeng's avatar
qiuqunfeng committed
617 618 619 620
		listenerStatusList = append(listenerStatusList, &GatewayListener{
			GatewayName: req.GatewayName,
			Namespace:   req.Namespace,
			RegionCode:  req.RegionCode,
621 622
			Port:        listenerPortInt,
			Hosts:       listener.Spec.HostNames,
qiuqunfeng's avatar
qiuqunfeng committed
623 624 625
		})
	}

626 627 628 629 630 631 632 633 634 635
	// for _, port := range portList {
	// 	listenerStatusList = append(listenerStatusList, &GatewayListener{
	// 		GatewayName: req.GatewayName,
	// 		Namespace:   req.Namespace,
	// 		RegionCode:  req.RegionCode,
	// 		Port:        port,
	// 		Enable:      true,
	// 	})
	// }

qiuqunfeng's avatar
qiuqunfeng committed
636 637
	return listenerStatusList, nil
}
638 639 640 641 642 643 644 645

func (s *wafService) EnableListenerWafs(ctx context.Context, req *EnableListenerWafsReq) error {

	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
		return fmt.Errorf("failed to get cluster client")
	}

646
	listenerList, err := client.Versioned.WafV1alpha1().Services(req.Namespace).List(ctx, metav1.ListOptions{LabelSelector: fmt.Sprintf("apigateway_name=%s", req.GatewayName)})
647
	if err != nil {
648
		log.Error().Msgf("failed to get listener list: %v", err)
649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667
		return err
	}

	portList := []int{}
	for _, listener := range listenerList.Items {
		n := strings.LastIndex(listener.Name, "-")
		if n == -1 {
			return fmt.Errorf("failed to get listener port: %v", listener.Name)
		}
		listenerPort := listener.Name[n+1:]
		listenerPortInt, err := strconv.Atoi(listenerPort)
		if err != nil {
			return fmt.Errorf("failed to parse listener port: %v", err)
		}
		portList = append(portList, listenerPortInt)
	}
	currentPortSet := sets.NewInt(portList...)

	desiredPortSet := sets.NewInt()
668
	wafMap := map[int]ListenerWaf{}
669
	for _, listener := range req.Listeners {
670
		// get port from listener.HostsAndPort, like hosts1@127.0.0.1@abc.com-8080
671
		index := strings.LastIndex(listener.HostsAndPort, "-")
672 673 674
		if index == -1 {
			return fmt.Errorf("failed to get listener port: %v", listener)
		}
675
		port := listener.HostsAndPort[index+1:]
676 677 678 679 680
		portInt, err := strconv.Atoi(port)
		if err != nil {
			return fmt.Errorf("failed to parse listener port: %v", err)
		}
		desiredPortSet.Insert(portInt)
681
		log.Info().Msgf("listener: %v", listener.Name)
682

683
		hosts := strings.Split(listener.HostsAndPort[:index], "@")
684 685 686 687 688
		wafMap[portInt] = ListenerWaf{
			Hosts:        hosts,
			HostsAndPort: listener.HostsAndPort,
			Name:         listener.Name,
		}
689 690 691 692
	}

	// enable WAF for ports that are in the desired port set but not in the current port set
	addingPortSet := desiredPortSet.Difference(currentPortSet)
693 694 695 696 697 698 699 700 701 702 703 704

	// Get mode from waf_services table
	wafService := &model.WafService{}
	err = s.db.Model(&model.WafService{}).Where("gateway_name = ?", req.GatewayName).First(wafService).Error
	if err != nil {
		if err == gorm.ErrRecordNotFound {
			return fmt.Errorf("waf service not found for gateway %s", req.GatewayName)
		}
		return fmt.Errorf("failed to query waf service: %v", err)
	}
	mode := WafMode(wafService.Mode)

705 706 707 708 709 710 711
	for _, port := range addingPortSet.List() {
		err := s.EnableListenerWaf(ctx, &EnableListenerWafReq{
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
712 713 714 715 716
			Port:         port,
			Hosts:        wafMap[port].Hosts,
			Enable:       true,
			Mode:         mode,
			ListenerName: wafMap[port].Name,
717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739
		})
		if err != nil {
			return fmt.Errorf("failed to enable listener WAF: %v", err)
		}
	}

	// delete WAF for ports that are not in the desired port set
	deletingPortSet := currentPortSet.Difference(desiredPortSet)
	for _, port := range deletingPortSet.List() {
		err := s.DeleteListenerWaf(ctx, &DeleteListenerReq{
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
			Port: port,
		})
		if err != nil {
			return fmt.Errorf("failed to delete listener WAF: %v", err)
		}
	}
	return nil
}
740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798

func generateListToken(sort []interface{}) (string, error) {
	var token []string

	for _, v := range sort {
		switch v.(type) {
		case string:
			token = append(token, v.(string))
		case float64:
			token = append(token, strconv.FormatInt(int64(v.(float64)), 10))
		default:
			// json-iterator 和 encoding/json 的Decode,返回的都是 encoding/json 的 Number
			t, ok := jsoniter.CastJsonNumber(v)
			if !ok {
				return "", fmt.Errorf("unsupported sort field. value: %v, type: %T", v, v)
			}
			token = append(token, t)
		}
	}

	return strings.Join(token, ","), nil
}

func (s *wafService) ListAttackLogs(ctx context.Context, req *AttackLogFilter) ([]AttackLog, string, error) {
	boolQuery := elastic.NewBoolQuery()
	if req.ServiceId != 0 {
		boolQuery.Must(elastic.NewTermQuery("service_id", req.ServiceId))
	}
	if req.Cluster != "" {
		boolQuery.Filter(elastic.NewTermQuery("cluster_key", req.Cluster))
	}
	if req.AttackUrl != "" {
		boolQuery.Filter(elastic.NewMatchPhraseQuery("attacked_url", req.AttackUrl).Slop(0))
	}
	if req.AttackIp != "" {
		boolQuery.Filter(elastic.NewMatchPhraseQuery("attack_ip", req.AttackIp).Slop(0))
	}
	if req.AttackApp != "" {
		boolQuery.Filter(elastic.NewMatchPhraseQuery("attacked_app", req.AttackApp).Slop(0))
	}
	if req.AttackType != "" {
		attackTypes := strings.Split(req.AttackType, ",")
		var terms []any
		for _, attackType := range attackTypes {
			terms = append(terms, attackType)
		}
		boolQuery.Filter(elastic.NewTermsQuery("attack_type", terms...))
	}
	if req.Action != "" {
		actions := strings.Split(req.Action, ",")
		var terms []interface{}
		for _, action := range actions {
			terms = append(terms, action)
		}
		boolQuery.Filter(elastic.NewTermsQuery("action", terms...))
	} else {
		boolQuery.Filter(elastic.NewBoolQuery().MustNot(elastic.NewTermQuery("action", "pass")))
	}

799 800
	hasStart := req.StartTime > 0
	hasEnd := req.EndTime > 0
801 802 803 804 805 806 807 808 809 810 811 812
	if hasStart || hasEnd {
		rangeQuery := elastic.NewRangeQuery("attack_time")
		if hasStart {
			rangeQuery.Gte(req.StartTime)
		}
		if hasEnd {
			rangeQuery.Lte(req.EndTime)
		}

		boolQuery.Filter(rangeQuery)
	}
	src, _ := boolQuery.Source()
813
	log.Info().Interface("src", src.(map[string]interface{})).Msg("find waf detections src")
814 815 816 817 818 819 820 821

	ss := s.elasticClient.Search("waf-detections*")
	if req.Token != "" {
		for _, t := range strings.Split(req.Token, ",") {
			ss.SearchAfter(t)
		}
	}

qiuqunfeng's avatar
debug  
qiuqunfeng committed
822
	log.Info().Interface("limit", req.Limit).Msg("limit")
823
	result, err := ss.Query(boolQuery).Size(req.Limit).
824 825 826 827 828 829 830 831
		SortBy(elastic.NewFieldSort("attack_time").Order(false),
			elastic.NewFieldSort("id.digit").Order(false)).
		Do(ctx)

	if err != nil {
		return nil, "", fmt.Errorf("failed to search waf detections: %v", err)
	}

832 833
	list := make([]model.WafDetection, len(result.Hits.Hits))
	endIdx := len(result.Hits.Hits) - 1
834
	pageToken := ""
835 836
	log.Info().Interface("res", result).Msg("list attack logs res")
	for i, hit := range result.Hits.Hits {
837
		log.Info().Interface("hit source", hit.Source).Msg("hit")
838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854
		wafDetection := model.WafDetection{}
		if err = json.Unmarshal(hit.Source, &wafDetection); err != nil {
			return nil, "", fmt.Errorf("failed to unmarshal waf detection: %v", err)
		}

		list[i] = wafDetection
		if i == endIdx {
			pageToken, err = generateListToken(hit.Sort)
			if err != nil {
				return nil, "", fmt.Errorf("failed to generate list token: %v", err)
			}
		}
	}

	attackLogs := make([]AttackLog, len(list))
	for i, wafDetection := range list {
		attackLogs[i] = AttackLog{
855
			Uuid:         wafDetection.ID,
856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880
			AttackTime:   wafDetection.AttackTime,
			AttackIp:     wafDetection.AttackIP,
			AttackedApp:  wafDetection.AttackedApp,
			AttackType:   wafDetection.AttackType,
			Action:       wafDetection.Action,
			ClusterKey:   wafDetection.ClusterKey,
			AttackedAddr: wafDetection.AttackedURL,
		}
	}
	return attackLogs, pageToken, nil
}

func (s *wafService) ListRules(ctx context.Context, regionCode, namespace, gatewayName, language, name string) ([]RuleGroupResp, error) {
	ruleCategories := []model.WafRuleCategory{}
	err := s.db.Model(&model.WafRuleCategory{}).Find(&ruleCategories).Error
	if err != nil {
		return nil, fmt.Errorf("failed to get waf service: %v", err)
	}

	ruleGroupResp := []RuleGroupResp{}
	wafService := &model.WafService{}
	err = s.db.Model(&model.WafService{}).Where("gateway_name = ? and namespace = ? and region_code = ?", gatewayName, namespace, regionCode).First(wafService).Error
	if err != nil {
		return nil, fmt.Errorf("failed to get waf service: %v", err)
	}
881
	if wafService.RuleCategoryStatus != nil && wafService.RuleCategoryStatus.Status == 1 {
882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903
		for _, category := range ruleCategories {
			for _, categoryID := range wafService.RuleCategoryStatus.CategoryID {
				if category.CategoryID == categoryID {
					category.Status = 1
				}
			}
			if language == "en" {
				ruleGroupResp = append(ruleGroupResp, RuleGroupResp{
					CategoryID:  category.CategoryID,
					Status:      category.Status,
					Category:    category.CategoryEN,
					Description: category.DescriptionEN,
				})
			} else {
				ruleGroupResp = append(ruleGroupResp, RuleGroupResp{
					CategoryID:  category.CategoryID,
					Status:      category.Status,
					Category:    category.CategoryZH,
					Description: category.DescriptionZH,
				})
			}
		}
904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921
	} else {
		for _, category := range ruleCategories {
			if language == "en" {
				ruleGroupResp = append(ruleGroupResp, RuleGroupResp{
					CategoryID:  category.CategoryID,
					Status:      category.Status,
					Category:    category.CategoryEN,
					Description: category.DescriptionEN,
				})
			} else {
				ruleGroupResp = append(ruleGroupResp, RuleGroupResp{
					CategoryID:  category.CategoryID,
					Status:      category.Status,
					Category:    category.CategoryZH,
					Description: category.DescriptionZH,
				})
			}
		}
922 923 924
	}
	return ruleGroupResp, nil
}
925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993 994 995 996 997 998 999 1000

func (s *wafService) getWafServiceMap(ctx context.Context, req *MatcherExpr) (map[string][]model.WafService, error) {
	svcMap := make(map[string][]model.WafService)
	wafServices := []model.WafService{}
	if req.Global {
		if err := s.db.WithContext(ctx).Model(&model.WafService{}).Find(&wafServices).Error; err != nil {
			return nil, err
		}
	} else {
		if err := s.db.WithContext(ctx).Model(&model.WafService{}).Where("id in ?", req.Scope).Find(&wafServices).Error; err != nil {
			return nil, err
		}
	}
	for _, wafService := range wafServices {
		svcMap[wafService.RegionCode] = append(svcMap[wafService.RegionCode], wafService)
	}
	return svcMap, nil
}

func (s *wafService) createConfigMap(ctx context.Context, req *MatcherExpr, regionCode string, wafSvc []model.WafService) error {
	client := s.clusterClientManager.GetClient(regionCode)
	if client == nil {
		return fmt.Errorf("failed to get cluster client")
	}

	gatewayNames := []string{}
	for _, wafSvc := range wafSvc {
		gatewayNames = append(gatewayNames, wafSvc.GatewayName)
	}
	scope := strings.Join(gatewayNames, ",")
	matchExpr := v1alpha1.MatchExpression{
		ID:     req.ID,
		Name:   req.Name,
		Scope:  scope,
		Mode:   req.Mode,
		Expr:   req.Expr,
		Status: req.Status,
	}
	matchExprJson, err := json.Marshal(matchExpr)
	if err != nil {
		return fmt.Errorf("failed to marshal match expression: %v", err)
	}
	name := fmt.Sprintf("waf-black-white-list-%d", req.ID)
	_, err = client.Clientset.CoreV1().ConfigMaps(wafSvc[0].Namespace).Create(ctx, &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
		},
		Data: map[string]string{
			"match-expression": string(matchExprJson),
		},
	}, metav1.CreateOptions{})
	if err != nil {
		return fmt.Errorf("failed to create config map: %v", err)
	}
	return nil
}

func (s *wafService) CreateBlackWhiteList(ctx context.Context, req *MatcherExpr) error {
	tErr := s.db.Transaction(func(tx *gorm.DB) error {
		matcherExpr := model.MatcherExpr{
			Name:   req.Name,
			Scope:  req.Scope,
			Mode:   req.Mode,
			Expr:   req.Expr,
			Global: req.Global,
		}
		err := s.db.WithContext(ctx).Create(&matcherExpr).Error
		if err != nil {
			return err
		}

		req.ID = matcherExpr.ID
		svcMap, err := s.getWafServiceMap(ctx, req)
		if err != nil {
			return err
		}
1001
		log.Info().Interface("svcMap", svcMap).Msg("svcMap")
1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028
		for regionCode, wafServices := range svcMap {
			err := s.createConfigMap(ctx, req, regionCode, wafServices)
			if err != nil {
				return err
			}
		}

		return nil
	})
	return tErr
}

func (s *wafService) UpdateBlackWhiteList(ctx context.Context, req *MatcherExpr) error {
	return nil
}

func (s *wafService) EnableBlackWhiteList(ctx context.Context, req *MatcherExpr) error {
	return nil
}

func (s *wafService) DeleteBlackWhiteList(ctx context.Context, req *MatcherExpr) error {
	return nil
}

func (s *wafService) GetBlackWhiteLists(ctx context.Context, req *MatcherExpr) ([]MatcherExpr, error) {
	return nil, nil
}