1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
package proxy
import (
"errors"
"sync"
"time"
"github.com/jumpserver/koko/pkg/i18n"
"github.com/jumpserver/koko/pkg/logger"
"github.com/jumpserver/koko/pkg/model"
"github.com/jumpserver/koko/pkg/service"
"github.com/jumpserver/koko/pkg/utils"
)
var sessionMap = make(map[string]*SwitchSession)
var lock = new(sync.RWMutex)
func HandleSessionTask(task model.TerminalTask) {
switch task.Name {
case "kill_session":
if ok := KillSession(task.Args); ok {
service.FinishTask(task.ID)
}
default:
}
}
func KillSession(sessionID string) bool {
lock.RLock()
defer lock.RUnlock()
if sw, ok := sessionMap[sessionID]; ok {
sw.Terminate()
return true
}
return false
}
func GetAliveSessions() []string {
lock.RLock()
defer lock.RUnlock()
sids := make([]string, 0, len(sessionMap))
for sid := range sessionMap {
sids = append(sids, sid)
}
return sids
}
func RemoveSession(sw *SwitchSession) {
lock.Lock()
defer lock.Unlock()
delete(sessionMap, sw.ID)
finishSession(sw)
}
func AddSession(sw *SwitchSession) {
lock.Lock()
defer lock.Unlock()
sessionMap[sw.ID] = sw
}
func CreateSession(p *ProxyServer) (sw *SwitchSession, err error) {
// 创建Session
sw = NewSwitchSession(p)
// Post到Api端
ok := postSession(sw)
msg := i18n.T("Connect with api server failed")
if !ok {
msg = utils.WrapperWarn(msg)
utils.IgnoreErrWriteString(p.UserConn, msg)
logger.Error(msg)
return sw, errors.New("connect api server failed")
}
// 获取系统用户的过滤规则，并设置
cmdRules, err := service.GetSystemUserFilterRules(p.SystemUser.ID)
if err != nil {
msg = utils.WrapperWarn(msg)
utils.IgnoreErrWriteString(p.UserConn, msg)
logger.Error(msg + err.Error())
return sw, errors.New("connect api server failed")
}
sw.SetFilterRules(cmdRules)
AddSession(sw)
return
}
func postSession(s *SwitchSession) bool {
data := s.MapData()
for i := 0; i < 5; i++ {
if service.CreateSession(data) {
return true
}
time.Sleep(200 * time.Millisecond)
}
return false
}
func finishSession(s *SwitchSession) {
data := s.MapData()
service.FinishSession(data)
logger.Debugf("Session %s has finished", s.ID)
}