117 lines
1.9 KiB
Go
117 lines
1.9 KiB
Go
package rand
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/rand"
|
|
"errors"
|
|
)
|
|
|
|
// define a flag that generates a random string
|
|
const (
|
|
Ldigit = 1 << iota
|
|
LlowerCase
|
|
LupperCase
|
|
LlowerAndUpperCase = LlowerCase | LupperCase
|
|
LdigitAndLowerCase = Ldigit | LlowerCase
|
|
LdigitAndUpperCase = Ldigit | LupperCase
|
|
LdigitAndLetter = Ldigit | LlowerCase | LupperCase
|
|
)
|
|
|
|
var (
|
|
digits = []byte("0123456789")
|
|
lowerCaseLetters = []byte("abcdefghijklmnopqrstuvwxyz")
|
|
upperCaseLetters = []byte("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
|
|
)
|
|
|
|
// definition error
|
|
var (
|
|
ErrInvalidFlag = errors.New("Invalid flag")
|
|
)
|
|
|
|
// Random generate a random string specifying the length of the random number
|
|
// and the random flag
|
|
func Random(length, flag int) (string, error) {
|
|
if length < 1 {
|
|
length = 6
|
|
}
|
|
|
|
source, err := getFlagSource(flag)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
b, err := randomBytesMod(length, byte(len(source)))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
var buf bytes.Buffer
|
|
for _, c := range b {
|
|
buf.WriteByte(source[c])
|
|
}
|
|
|
|
return buf.String(), nil
|
|
}
|
|
|
|
func getFlagSource(flag int) ([]byte, error) {
|
|
var source []byte
|
|
|
|
if flag&Ldigit > 0 {
|
|
source = append(source, digits...)
|
|
}
|
|
|
|
if flag&LlowerCase > 0 {
|
|
source = append(source, lowerCaseLetters...)
|
|
}
|
|
|
|
if flag&LupperCase > 0 {
|
|
source = append(source, upperCaseLetters...)
|
|
}
|
|
|
|
sourceLen := len(source)
|
|
if sourceLen == 0 {
|
|
return nil, ErrInvalidFlag
|
|
}
|
|
return source, nil
|
|
}
|
|
|
|
func randomBytesMod(length int, mod byte) ([]byte, error) {
|
|
b := make([]byte, length)
|
|
max := 255 - 255%mod
|
|
i := 0
|
|
|
|
LROOT:
|
|
for {
|
|
r, err := randomBytes(length + length/4)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
for _, c := range r {
|
|
if c >= max {
|
|
// Skip this number to avoid modulo bias
|
|
continue
|
|
}
|
|
|
|
b[i] = c % mod
|
|
i++
|
|
if i == length {
|
|
break LROOT
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
return b, nil
|
|
}
|
|
|
|
func randomBytes(length int) ([]byte, error) {
|
|
b := make([]byte, length)
|
|
_, err := rand.Read(b)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return b, nil
|
|
}
|