github.com/TeaOSLab/EdgeNode@v1.3.8/internal/waf/captcha_generator.go (about)

     1  // Copyright 2023 GoEdge CDN goedge.cdn@gmail.com. All rights reserved. Official site: https://goedge.cn .
     2  
     3  package waf
     4  
     5  import (
     6  	"bytes"
     7  	"github.com/dchest/captcha"
     8  	"github.com/iwind/TeaGo/rands"
     9  	"io"
    10  	"time"
    11  )
    12  
    13  // CaptchaGenerator captcha generator
    14  type CaptchaGenerator struct {
    15  	store captcha.Store
    16  }
    17  
    18  func NewCaptchaGenerator() *CaptchaGenerator {
    19  	return &CaptchaGenerator{
    20  		store: captcha.NewMemoryStore(100_000, 5*time.Minute),
    21  	}
    22  }
    23  
    24  // NewCaptcha create new captcha
    25  func (this *CaptchaGenerator) NewCaptcha(length int) (captchaId string) {
    26  	captchaId = rands.HexString(16)
    27  
    28  	if length <= 0 || length > 20 {
    29  		length = 4
    30  	}
    31  
    32  	this.store.Set(captchaId, captcha.RandomDigits(length))
    33  	return
    34  }
    35  
    36  // WriteImage write image to front writer
    37  func (this *CaptchaGenerator) WriteImage(w io.Writer, id string, width, height int) error {
    38  	var d = this.store.Get(id, false)
    39  	if d == nil {
    40  		return captcha.ErrNotFound
    41  	}
    42  	_, err := captcha.NewImage(id, d, width, height).WriteTo(w)
    43  	return err
    44  }
    45  
    46  // Verify user input
    47  func (this *CaptchaGenerator) Verify(id string, digits string) bool {
    48  	var countDigits = len(digits)
    49  	if countDigits == 0 {
    50  		return false
    51  	}
    52  	var value = this.store.Get(id, true)
    53  	if len(value) != countDigits {
    54  		return false
    55  	}
    56  
    57  	var nb = make([]byte, countDigits)
    58  	for i := 0; i < countDigits; i++ {
    59  		var d = digits[i]
    60  		if d >= '0' && d <= '9' {
    61  			nb[i] = d - '0'
    62  		}
    63  	}
    64  
    65  	return bytes.Equal(nb, value)
    66  }
    67  
    68  // Get captcha data
    69  func (this *CaptchaGenerator) Get(id string) []byte {
    70  	return this.store.Get(id, false)
    71  }