Skip to content

K
koko
Commit cb4f6cd4 authored by Eric's avatar Eric
Browse files
Options

[Update] fix some elasticsearch bugs

parent 0e374565
Hide whitespace changes
Inline Side-by-side
Showing with 72 additions and 63 deletions
pkg/config/config.go
View file @ cb4f6cd4
......@@ -12,36 +12,36 @@ import (
)
type Config struct {
AssetListPageSize string `json:"TERMINAL_ASSET_LIST_PAGE_SIZE"`
AssetListSortBy string `json:"TERMINAL_ASSET_LIST_SORT_BY"`
HeaderTitle string `json:"TERMINAL_HEADER_TITLE"`
HostKey string `json:"TERMINAL_HOST_KEY" yaml:"HOST_KEY"`
PasswordAuth bool `json:"TERMINAL_PASSWORD_AUTH" yaml:"PASSWORD_AUTH"`
PublicKeyAuth bool `json:"TERMINAL_PUBLIC_KEY_AUTH" yaml:"PUBLIC_KEY_AUTH"`
CommandStorage map[string]string `json:"TERMINAL_COMMAND_STORAGE"`
ReplayStorage map[string]string `json:"TERMINAL_REPLAY_STORAGE" yaml:"REPLAY_STORAGE"`
SessionKeepDuration int `json:"TERMINAL_SESSION_KEEP_DURATION"`
TelnetRegex string `json:"TERMINAL_TELNET_REGEX"`
MaxIdleTime int `json:"SECURITY_MAX_IDLE_TIME"`
SftpRoot string `json:"TERMINAL_SFTP_ROOT" yaml:"SFTP_ROOT"`
Name string `yaml:"NAME"`
SecretKey string `yaml:"SECRET_KEY"`
HostKeyFile string `yaml:"HOST_KEY_FILE"`
CoreHost string `yaml:"CORE_HOST"`
BootstrapToken string `yaml:"BOOTSTRAP_TOKEN"`
BindHost string `yaml:"BIND_HOST"`
SSHPort int `yaml:"SSHD_PORT"`
HTTPPort int `yaml:"HTTPD_PORT"`
SSHTimeout int `yaml:"SSH_TIMEOUT"`
AccessKey string `yaml:"ACCESS_KEY"`
AccessKeyFile string `yaml:"ACCESS_KEY_FILE"`
LogLevel string `yaml:"LOG_LEVEL"`
HeartbeatDuration int `yaml:"HEARTBEAT_INTERVAL"`
RootPath string `yaml:"ROOT_PATH"`
Comment string `yaml:"COMMENT"`
Language string `yaml:"LANG"`
LanguageCode string `yaml:"LANGUAGE_CODE"` // Abandon
UploadFailedReplay bool `yaml:"UPLOAD_FAILED_REPLAY_ON_START"`
AssetListPageSize string `json:"TERMINAL_ASSET_LIST_PAGE_SIZE"`
AssetListSortBy string `json:"TERMINAL_ASSET_LIST_SORT_BY"`
HeaderTitle string `json:"TERMINAL_HEADER_TITLE"`
HostKey string `json:"TERMINAL_HOST_KEY" yaml:"HOST_KEY"`
PasswordAuth bool `json:"TERMINAL_PASSWORD_AUTH" yaml:"PASSWORD_AUTH"`
PublicKeyAuth bool `json:"TERMINAL_PUBLIC_KEY_AUTH" yaml:"PUBLIC_KEY_AUTH"`
CommandStorage map[string]interface{} `json:"TERMINAL_COMMAND_STORAGE"`
ReplayStorage map[string]interface{} `json:"TERMINAL_REPLAY_STORAGE" yaml:"REPLAY_STORAGE"`
SessionKeepDuration int `json:"TERMINAL_SESSION_KEEP_DURATION"`
TelnetRegex string `json:"TERMINAL_TELNET_REGEX"`
MaxIdleTime int `json:"SECURITY_MAX_IDLE_TIME"`
SftpRoot string `json:"TERMINAL_SFTP_ROOT" yaml:"SFTP_ROOT"`
Name string `yaml:"NAME"`
SecretKey string `yaml:"SECRET_KEY"`
HostKeyFile string `yaml:"HOST_KEY_FILE"`
CoreHost string `yaml:"CORE_HOST"`
BootstrapToken string `yaml:"BOOTSTRAP_TOKEN"`
BindHost string `yaml:"BIND_HOST"`
SSHPort int `yaml:"SSHD_PORT"`
HTTPPort int `yaml:"HTTPD_PORT"`
SSHTimeout int `yaml:"SSH_TIMEOUT"`
AccessKey string `yaml:"ACCESS_KEY"`
AccessKeyFile string `yaml:"ACCESS_KEY_FILE"`
LogLevel string `yaml:"LOG_LEVEL"`
HeartbeatDuration int `yaml:"HEARTBEAT_INTERVAL"`
RootPath string `yaml:"ROOT_PATH"`
Comment string `yaml:"COMMENT"`
Language string `yaml:"LANG"`
LanguageCode string `yaml:"LANGUAGE_CODE"` // Abandon
UploadFailedReplay bool `yaml:"UPLOAD_FAILED_REPLAY_ON_START"`
}
func (c *Config) EnsureConfigValid() {
......@@ -122,8 +122,8 @@ var Conf = &Config{
RootPath: rootPath,
Comment: "Coco",
Language: "zh",
ReplayStorage: map[string]string{"TYPE": "server"},
CommandStorage: map[string]string{"TYPE": "server"},
ReplayStorage: map[string]interface{}{"TYPE": "server"},
CommandStorage: map[string]interface{}{"TYPE": "server"},
UploadFailedReplay: true,
}
......
pkg/proxy/recorderstorage.go
View file @ cb4f6cd4
......@@ -43,34 +43,34 @@ func NewReplayStorage() ReplayStorage {
}
switch tp {
case "azure":
endpointSuffix := cf["ENDPOINT_SUFFIX"]
endpointSuffix := cf["ENDPOINT_SUFFIX"].(string)
if endpointSuffix == "" {
endpointSuffix = "core.chinacloudapi.cn"
}
return &AzureReplayStorage{
accountName: cf["ACCOUNT_NAME"],
accountKey: cf["ACCOUNT_KEY"],
containerName: cf["CONTAINER_NAME"],
accountName: cf["ACCOUNT_NAME"].(string),
accountKey: cf["ACCOUNT_KEY"].(string),
containerName: cf["CONTAINER_NAME"].(string),
endpointSuffix: endpointSuffix,
}
case "oss":
return &OSSReplayStorage{
endpoint: cf["ENDPOINT"],
bucket: cf["BUCKET"],
accessKey: cf["ACCESS_KEY"],
secretKey: cf["SECRET_KEY"],
endpoint: cf["ENDPOINT"].(string),
bucket: cf["BUCKET"].(string),
accessKey: cf["ACCESS_KEY"].(string),
secretKey: cf["SECRET_KEY"].(string),
}
case "s3":
bucket := cf["BUCKET"]
bucket := cf["BUCKET"].(string)
if bucket == "" {
bucket = "jumpserver"
}
return &S3ReplayStorage{
bucket: bucket,
region: cf["REGION"],
accessKey: cf["ACCESS_KEY"],
secretKey: cf["SECRET_KEY"],
endpoint: cf["ENDPOINT"],
region: cf["REGION"].(string),
accessKey: cf["ACCESS_KEY"].(string),
secretKey: cf["SECRET_KEY"].(string),
endpoint: cf["ENDPOINT"].(string),
}
default:
return defaultReplayStorage
......@@ -85,17 +85,19 @@ func NewCommandStorage() CommandStorage {
}
switch tp {
case "es", "elasticsearch":
hosts := cf["HOSTS"]
index := cf["INDEX"]
docType := cf["DOC_TYPE"]
hostsArray := strings.Split(strings.Trim(hosts, ","), ",")
var hosts = make([]string, len(cf["HOSTS"].([]interface{})))
for i, item := range cf["HOSTS"].([]interface{}) {
hosts[i] = item.(string)
}
index := cf["INDEX"].(string)
docType := cf["DOC_TYPE"].(string)
if index == "" {
index = "jumpserver"
}
if docType == "" {
docType = "command_store"
}
return &ESCommandStorage{hosts: hostsArray, index: index, docType: docType}
return &ESCommandStorage{hosts: hosts, index: index, docType: docType}
default:
return defaultCommandStorage
}
......@@ -115,11 +117,7 @@ type ESCommandStorage struct {
}
func (es *ESCommandStorage) BulkSave(commands []*model.Command) (err error) {
data, err := json.Marshal(commands)
if err != nil {
return
}
var buf bytes.Buffer
esClinet, err := elasticsearch.NewClient(elasticsearch.Config{
Addresses: es.hosts,
})
......@@ -127,15 +125,25 @@ func (es *ESCommandStorage) BulkSave(commands []*model.Command) (err error) {
logger.Error(err.Error())
return
}
for _, item := range commands {
meta := []byte(fmt.Sprintf(`{ "index" : { } }%s`, "\n"))
data, err := json.Marshal(item)
if err != nil {
return err
}
data = append(data, "\n"...)
buf.Write(meta)
buf.Write(data)
}
_, err = esClinet.Bulk(bytes.NewBuffer(data),
_, err = esClinet.Bulk(bytes.NewReader(buf.Bytes()),
esClinet.Bulk.WithIndex(es.index), esClinet.Bulk.WithDocumentType(es.docType))
if err == nil {
logger.Debug("Successfully uploaded total %d commands to Elasticsearch\n", len(commands))
if err != nil {
logger.Error(err.Error())
}
return
}
func NewFileCommandStorage(name string) (storage *FileCommandStorage, err error) {
file, err := os.Create(name)
if err != nil {
......@@ -181,6 +189,7 @@ func (o *OSSReplayStorage) Upload(gZipFilePath, target string) (err error) {
}
bucket, err := client.Bucket(o.bucket)
if err != nil {
logger.Error(err.Error())
return
}
return bucket.PutObjectFromFile(target, gZipFilePath)
......@@ -217,8 +226,8 @@ func (s *S3ReplayStorage) Upload(gZipFilePath, target string) (err error) {
Key: aws.String(target),
Body: file,
})
if err == nil {
logger.Debug("Successfully uploaded %q to %q\n", file.Name(), s.bucket)
if err != nil {
logger.Error(err.Error())
}
return
......@@ -250,8 +259,8 @@ func (a *AzureReplayStorage) Upload(gZipFilePath, target string) (err error) {
_, err = azblob.UploadFileToBlockBlob(context.TODO(), file, blobURL, azblob.UploadToBlockBlobOptions{
BlockSize: 4 * 1024 * 1024,
Parallelism: 16})
if err == nil {
logger.Debug("Successfully uploaded %q to Azure\n", file.Name())
if err != nil {
logger.Error(err.Error())
}
return
}
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment