waf.go 28.6 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 24
	"gitlab.com/tensorsecurity-rd/waf-console/pkg/apis/waf.security.io/v1alpha1"
	"gorm.io/gorm"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
25
	"k8s.io/apimachinery/pkg/util/sets"
qiuqunfeng's avatar
qiuqunfeng committed
26 27 28
)

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

35 36
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
37 38
}

qiuqunfeng's avatar
qiuqunfeng committed
39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
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
57
	}
58 59 60 61 62 63 64 65 66 67 68 69 70 71
	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
72 73 74 75 76
	return &WafService{
		GatewayName: wafService.GatewayName,
		Mode:        wafService.Mode,
		RuleNum:     wafService.RuleNum,
		AttackNum:   wafService.AttackNum,
77
		Listeners:   listeners,
qiuqunfeng's avatar
qiuqunfeng committed
78
	}, nil
qiuqunfeng's avatar
qiuqunfeng committed
79 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
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
130 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
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
192 193 194
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
195 196
	service := &v1alpha1.Service{
		ObjectMeta: metav1.ObjectMeta{
qiuqunfeng's avatar
qiuqunfeng committed
197
			Name:      name,
qiuqunfeng's avatar
qiuqunfeng committed
198
			Namespace: req.Namespace,
qiuqunfeng's avatar
qiuqunfeng committed
199 200 201
			Labels: map[string]string{
				"apigateway_name": req.GatewayName,
			},
qiuqunfeng's avatar
qiuqunfeng committed
202 203 204
		},
		Spec: v1alpha1.ServiceSpec{
			HostNames:   req.Host,
205
			ServiceName: req.ListenerName,
qiuqunfeng's avatar
qiuqunfeng committed
206 207
			Port:        req.Port,
			Workload: v1alpha1.WorkloadRef{
208 209 210 211
				Kind:       req.GatewayName,
				Name:       req.ListenerName,
				Namespace:  req.Namespace,
				ClusterKey: req.RegionCode,
qiuqunfeng's avatar
qiuqunfeng committed
212
			},
qiuqunfeng's avatar
qiuqunfeng committed
213 214 215 216 217 218 219
			Uri: &v1alpha1.StringMatch{
				Prefix: "/",
			},
			LogConfig: &v1alpha1.LogConfig{
				Enable: 1,
				Level:  "info",
			},
220 221
			Mode:      string(req.Mode),
			ServiceID: req.ServiceID,
qiuqunfeng's avatar
qiuqunfeng committed
222 223
		},
	}
qiuqunfeng's avatar
qiuqunfeng committed
224

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

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

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

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

qiuqunfeng's avatar
qiuqunfeng committed
252 253 254
func (s *wafService) DeleteListenerWaf(ctx context.Context, req *DeleteListenerReq) error {
	client := s.clusterClientManager.GetClient(req.RegionCode)
	if client == nil {
255
		return fmt.Errorf("failed to get cluster client for region %s", req.RegionCode)
qiuqunfeng's avatar
qiuqunfeng committed
256 257 258 259 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
	}
	name := fmt.Sprintf("%s-%d", req.GatewayName, req.Port)
	if err := client.WafV1alpha1().Services(req.Namespace).Delete(ctx, name, metav1.DeleteOptions{}); err != nil {
		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)
		}
	}
289 290 291 292 293 294 295 296 297
	// 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)
	}
	listenerList, err := client.WafV1alpha1().Services(req.Namespace).List(ctx, metav1.ListOptions{LabelSelector: fmt.Sprintf("apigateway_name=%s", req.GatewayName)})
	if err != nil {
		return nil, fmt.Errorf("failed to get listener list: %v", err)
	}
298
	var wg sync.WaitGroup
299
	for _, listener := range listenerList.Items {
300
		wg.Add(1)
301
		listener := listener // Create new variable for goroutine
302
		listener.Spec.Mode = string(req.Mode)
303
		go func() {
304
			defer wg.Done()
305 306 307 308 309 310 311
			log.Info().Msgf("update WAF service mode: %v", listener.Name)
			_, err := client.WafV1alpha1().Services(req.Namespace).Update(ctx, &listener, metav1.UpdateOptions{})
			if err != nil {
				log.Error().Msgf("failed to update WAF service mode: %v", err)
			}
		}()
	}
312
	wg.Wait()
313 314 315 316
	return &WafService{
		GatewayName: req.GatewayName,
		Mode:        string(req.Mode),
	}, nil
qiuqunfeng's avatar
commit  
qiuqunfeng committed
317
}
qiuqunfeng's avatar
commit  
qiuqunfeng committed
318

319 320 321 322 323 324 325 326
// 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
327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347

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
348
	jsonFile, err := os.ReadFile("rules/waf-rules.json")
qiuqunfeng's avatar
commit  
qiuqunfeng committed
349 350 351 352
	if err != nil {
		return fmt.Errorf("error reading yaml file: %v", err)
	}

353 354
	// err = yaml.Unmarshal(yamlFile, &categories)
	err = json.Unmarshal(jsonFile, &categories)
qiuqunfeng's avatar
commit  
qiuqunfeng committed
355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381
	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,
			CategoryEN:    category.Catagory.EN,
			CategoryZH:    category.Catagory.Zh,
			DescriptionEN: category.Description.EN,
			DescriptionZH: category.Description.Zh,
			Rules:         model.RuleList(rules),
		}
qiuqunfeng's avatar
commit  
qiuqunfeng committed
382 383 384 385
		err = s.db.Table("waf_rule_categories").Create(&model).Error
		if err != nil {
			return err
		}
qiuqunfeng's avatar
commit  
qiuqunfeng committed
386 387 388 389
	}

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

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
}

406 407 408 409 410 411
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
412

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

func (s *wafService) EnableListenerWaf(ctx context.Context, req *EnableListenerWafReq) error {
417 418
	if req.Enable {
		log.Info().Msgf("Create WAF for listener %s", req.GatewayName)
419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437
		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
438 439 440 441 442
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
443 444 445 446
			Port:         uint32(req.Port),
			Host:         req.Hosts,
			Mode:         req.Mode,
			ListenerName: req.ListenerName,
447
			ServiceID:    serviceID,
qiuqunfeng's avatar
qiuqunfeng committed
448
		})
qiuqunfeng's avatar
commit  
qiuqunfeng committed
449 450 451
		if err != nil {
			return err
		}
qiuqunfeng's avatar
qiuqunfeng committed
452
	} else {
453 454
		log.Info().Msgf("Delete WAF for listener %s", req.GatewayName)
		err := s.DeleteListenerWaf(ctx, &DeleteListenerReq{
qiuqunfeng's avatar
qiuqunfeng committed
455 456 457 458 459 460 461
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
			Port: req.Port,
		})
qiuqunfeng's avatar
commit  
qiuqunfeng committed
462 463 464
		if err != nil {
			return err
		}
qiuqunfeng's avatar
qiuqunfeng committed
465 466 467 468
	}
	return nil
}

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

475
func (s *wafService) listListenerFromApiGateway(ctx context.Context, apiGatewayCrn string, regionCode string, cookie string) ([]GatewayRespListenerData, error) {
476 477 478 479 480 481 482
	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)
	}
483
	request, err := http.NewRequestWithContext(ctx, "POST", "https://csm.console.test.tg.unicom.local/apigatewaymng/listener/lf-tst7/list_listeners", bytes.NewBuffer(body))
484 485 486 487
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %v", err)
	}
	request.Header.Set("Cookie", cookie)
qiuqunfeng's avatar
commit  
qiuqunfeng committed
488 489 490 491 492 493 494 495
	// 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)
496 497 498 499 500
	if err != nil {
		return nil, fmt.Errorf("failed to get listener list: %v", err)
	}
	defer resp.Body.Close()

501
	log.Info().Msgf("resp: %v", resp)
502 503 504 505 506 507
	// 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)
	}
508
	log.Info().Msgf("response: %v", response)
509 510 511
	return response.Data, nil
}

qiuqunfeng's avatar
qiuqunfeng committed
512 513
func (s *wafService) EnableGatewayWaf(ctx context.Context, req *EnableGatewayWafReq) error {
	if req.Enable {
514
		listeners, err := s.listListenerFromApiGateway(ctx, req.ApiGatewayCrn, req.RegionCode, req.Cookie)
qiuqunfeng's avatar
qiuqunfeng committed
515 516 517
		if err != nil {
			return fmt.Errorf("failed to get listener list: %v", err)
		}
qiuqunfeng's avatar
commit  
qiuqunfeng committed
518
		log.Info().Msgf("listeners: %v", listeners)
qiuqunfeng's avatar
qiuqunfeng committed
519 520
		// Create WAF for each listener
		for _, listener := range listeners {
qiuqunfeng's avatar
commit  
qiuqunfeng committed
521 522
			gatewayName := getGatewayNameFromCrn(listener.ApiGatewayCrn)
			namespace := fmt.Sprintf("%s-%s", listener.CreateAccountName, listener.CreateAccountID)
qiuqunfeng's avatar
qiuqunfeng committed
523 524
			if _, err := s.CreateWaf(ctx, &CreateWafReq{
				GatewateInfo: GatewateInfo{
qiuqunfeng's avatar
commit  
qiuqunfeng committed
525 526
					GatewayName: gatewayName,
					Namespace:   namespace,
qiuqunfeng's avatar
qiuqunfeng committed
527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 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 598 599 600 601 602 603 604 605
					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)
	if err := client.WafV1alpha1().Services(req.Namespace).DeleteCollection(ctx, metav1.DeleteOptions{}, metav1.ListOptions{LabelSelector: labelSelector}); err != nil {
		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")
	}

	listenerList, err := client.WafV1alpha1().Services(req.Namespace).List(ctx, metav1.ListOptions{LabelSelector: fmt.Sprintf("apigateway_name=%s", req.GatewayName)})
	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
606
			return nil, fmt.Errorf("failed to get listener port: %v", listener.Name)
qiuqunfeng's avatar
qiuqunfeng committed
607 608 609 610
		}
		listenerPort := listener.Name[n+1:]
		listenerPortInt, err := strconv.Atoi(listenerPort)
		if err != nil {
qiuqunfeng's avatar
commit  
qiuqunfeng committed
611
			return nil, fmt.Errorf("failed to parse listener port: %v", err)
qiuqunfeng's avatar
qiuqunfeng committed
612 613
		}

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

625 626 627 628 629 630 631 632 633 634
	// 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
635 636
	return listenerStatusList, nil
}
637 638 639 640 641 642 643 644 645 646

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")
	}

	listenerList, err := client.WafV1alpha1().Services(req.Namespace).List(ctx, metav1.ListOptions{LabelSelector: fmt.Sprintf("apigateway_name=%s", req.GatewayName)})
	if err != nil {
647
		log.Error().Msgf("failed to get listener list: %v", err)
648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665 666
		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()
667
	wafMap := map[int]ListenerWaf{}
668
	for _, listener := range req.Listeners {
669
		// get port from listener.HostsAndPort, like hosts1@127.0.0.1@abc.com-8080
670
		index := strings.LastIndex(listener.HostsAndPort, "-")
671 672 673
		if index == -1 {
			return fmt.Errorf("failed to get listener port: %v", listener)
		}
674
		port := listener.HostsAndPort[index+1:]
675 676 677 678 679
		portInt, err := strconv.Atoi(port)
		if err != nil {
			return fmt.Errorf("failed to parse listener port: %v", err)
		}
		desiredPortSet.Insert(portInt)
680
		log.Info().Msgf("listener: %v", listener.Name)
681

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

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

	// 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)

704 705 706 707 708 709 710
	for _, port := range addingPortSet.List() {
		err := s.EnableListenerWaf(ctx, &EnableListenerWafReq{
			GatewateInfo: GatewateInfo{
				GatewayName: req.GatewayName,
				Namespace:   req.Namespace,
				RegionCode:  req.RegionCode,
			},
711 712 713 714 715
			Port:         port,
			Hosts:        wafMap[port].Hosts,
			Enable:       true,
			Mode:         mode,
			ListenerName: wafMap[port].Name,
716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738
		})
		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
}
739 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

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")))
	}

798 799
	hasStart := req.StartTime > 0
	hasEnd := req.EndTime > 0
800 801 802 803 804 805 806 807 808 809 810 811
	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()
812
	log.Info().Interface("src", src.(map[string]interface{})).Msg("find waf detections src")
813 814 815 816 817 818 819 820

	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
821
	log.Info().Interface("limit", req.Limit).Msg("limit")
822
	result, err := ss.Query(boolQuery).Size(req.Limit).
823 824 825 826 827 828 829 830
		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)
	}

831 832
	list := make([]model.WafDetection, len(result.Hits.Hits))
	endIdx := len(result.Hits.Hits) - 1
833
	pageToken := ""
834 835
	log.Info().Interface("res", result).Msg("list attack logs res")
	for i, hit := range result.Hits.Hits {
836
		log.Info().Interface("hit source", hit.Source).Msg("hit")
837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853
		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{
854
			Uuid:         wafDetection.ID,
855 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 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905
			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)
	}
	if wafService.RuleCategoryStatus.Status == 1 {
		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,
				})
			}
		}
	}
	return ruleGroupResp, nil
}