Commit 96ee0fb8 authored by ibuler's avatar ibuler

[Update] 修改一些不优雅的代码

No related merge requests found
......@@ -11,6 +11,7 @@ import (
"cocogo/pkg/config"
"cocogo/pkg/i18n"
"cocogo/pkg/logger"
"cocogo/pkg/utils"
)
var defaultTitle string
......@@ -30,6 +31,7 @@ func (mi *MenuItem) Text() string {
cm := ColorMeta{GreenBoldColor: "\033[1;32m", ColorEnd: "\033[0m"}
line := fmt.Sprintf(i18n.T("\t%d) Enter {{.GreenBoldColor}}%s{{.ColorEnd}} to %s.%s"), mi.id, mi.instruct, mi.helpText, "\r\n")
tmpl := template.Must(template.New("item").Parse(line))
var buf bytes.Buffer
err := tmpl.Execute(&buf, cm)
if err != nil {
......@@ -42,7 +44,7 @@ func (mi *MenuItem) Text() string {
type Menu []MenuItem
func init() {
defaultTitle = i18n.T("Welcome to use Jumpserver open source fortress system")
defaultTitle = utils.WrapperTitle(i18n.T("Welcome to use Jumpserver open source fortress system"))
menu = Menu{
{id: 1, instruct: "ID", helpText: i18n.T("directly login")},
{id: 2, instruct: i18n.T("part IP, Hostname, Comment"), helpText: i18n.T("to search login if unique")},
......@@ -50,9 +52,8 @@ func init() {
{id: 4, instruct: "p", helpText: i18n.T("display the host you have permission")},
{id: 5, instruct: "g", helpText: "display the node that you have permission"},
{id: 6, instruct: "r", helpText: "refresh your assets and nodes"},
{id: 7, instruct: "s", helpText: "switch Chinese-english language"},
{id: 8, instruct: "h", helpText: "print help"},
{id: 9, instruct: "q", helpText: "exit"},
{id: 7, instruct: "h", helpText: "print help"},
{id: 8, instruct: "q", helpText: "exit"},
}
}
......@@ -66,12 +67,16 @@ func displayBanner(sess ssh.Session, user string) {
if config.Conf.HeaderTitle != "" {
title = config.Conf.HeaderTitle
}
welcomeMsg := CharClear + CharTab + user + " " + title + CharNewLine
prefix := utils.CharClear + utils.CharTab + utils.CharTab
suffix := utils.CharNewLine + utils.CharNewLine
welcomeMsg := prefix + utils.WrapperTitle(user+",") + " " + title + suffix
_, err := io.WriteString(sess, welcomeMsg)
if err != nil {
logger.Error("Send to client error, %s", err)
return
}
for _, v := range menu {
_, _ = io.WriteString(sess, v.Text())
utils.IgnoreErrWriteString(sess, v.Text())
}
}
package handler
const (
GreenColorCode = "\033[32m"
CharClear = "\x1b[H\x1b[2J"
CharTab = "\t"
CharNewLine = "\r\n"
AssetsMapKey = "AssetsMapKey"
AssetNodesMapKey = "AssetNodesKey"
)
......@@ -29,12 +29,7 @@ func SessionHandler(sess ssh.Session) {
_, _, ptyOk := sess.Pty()
if ptyOk {
ctx, cancel := cctx.NewContext(sess)
fmt.Println(ctx.User())
handler := &InteractiveHandler{
sess: sess,
user: ctx.User(),
term: terminal.NewTerminal(sess, "Opt> "),
}
handler := newInteractiveHandler(sess, ctx.User())
logger.Infof("New connection from: %s %s", sess.User(), sess.RemoteAddr().String())
handler.Dispatch(ctx)
cancel()
......@@ -46,29 +41,34 @@ func SessionHandler(sess ssh.Session) {
}
}
func newInteractiveHandler(sess ssh.Session, user *model.User) *InteractiveHandler {
term := terminal.NewTerminal(sess, "Opt> ")
handler := &InteractiveHandler{sess: sess, user: user, term: term}
handler.Initial()
return handler
}
type InteractiveHandler struct {
sess ssh.Session
term *terminal.Terminal
user *model.User
sess ssh.Session
user *model.User
term *terminal.Terminal
assetSelect *model.Asset
systemUserSelect *model.SystemUser
assets model.AssetList
searchResult model.AssetList
nodes model.NodeList
onceLoad sync.Once
mu sync.RWMutex
mu *sync.RWMutex
}
func (h *InteractiveHandler) displayBanner() {
displayBanner(h.sess, h.user.Name)
func (h *InteractiveHandler) Initial() {
h.displayBanner()
h.loadUserAssets()
h.loadUserAssetNodes()
}
func (h *InteractiveHandler) preDispatch() {
h.displayBanner()
h.onceLoad.Do(func() {
h.loadUserAssets()
h.loadUserAssetNodes()
})
func (h *InteractiveHandler) displayBanner() {
displayBanner(h.sess, h.user.Name)
}
func (h *InteractiveHandler) watchWinSizeChange(winCh <-chan ssh.Window, done <-chan struct{}) {
......@@ -88,8 +88,6 @@ func (h *InteractiveHandler) watchWinSizeChange(winCh <-chan ssh.Window, done <-
}
func (h *InteractiveHandler) Dispatch(ctx cctx.Context) {
h.preDispatch()
fmt.Println(h.user)
_, winCh, _ := h.sess.Pty()
for {
doneChan := make(chan struct{})
......@@ -113,8 +111,6 @@ func (h *InteractiveHandler) Dispatch(ctx cctx.Context) {
h.Proxy(ctx)
case "g":
h.displayNodes(h.nodes)
case "s":
h.changeLanguage()
case "h":
h.displayBanner()
case "r":
......@@ -144,6 +140,10 @@ func (h *InteractiveHandler) Dispatch(ctx cctx.Context) {
continue
}
}
default:
assets := h.searchAsset(line)
h.searchResult = assets
h.displayAssetsOrProxy(assets)
}
}
}
......@@ -151,9 +151,9 @@ func (h *InteractiveHandler) Dispatch(ctx cctx.Context) {
func (h *InteractiveHandler) chooseSystemUser(systemUsers []model.SystemUser) model.SystemUser {
table := tablewriter.NewWriter(h.sess)
table.SetHeader([]string{"ID", "UserName"})
table.SetHeader([]string{"ID", "Username"})
for i := 0; i < len(systemUsers); i++ {
table.Append([]string{strconv.Itoa(i + 1), systemUsers[i].UserName})
table.Append([]string{strconv.Itoa(i + 1), systemUsers[i].Username})
}
table.SetBorder(false)
count := 0
......@@ -325,10 +325,12 @@ func (h *InteractiveHandler) searchNodeAssets(num int) (assets []model.Asset) {
}
func (h *InteractiveHandler) Proxy(ctx context.Context) {
h.assetSelect = &model.Asset{Hostname: "centos", Port: 32768, Ip: "127.0.0.1"}
h.systemUserSelect = &model.SystemUser{Name: "web", UserName: "root", Password: "screencast"}
h.assetSelect = &model.Asset{Hostname: "centos", Port: 22, Ip: "192.168.244.185", Protocol: "ssh"}
h.systemUserSelect = &model.SystemUser{Id: "5dd8b5a0-8cdb-4857-8629-faf811c525e1", Name: "web", Username: "root", Password: "redhat", Protocol: "telnet"}
userConn := &proxy.UserSSHConnection{Session: h.sess}
p := proxy.ProxyServer{
Session: h.sess,
UserConn: userConn,
User: h.user,
Asset: h.assetSelect,
SystemUser: h.systemUserSelect,
......@@ -348,7 +350,7 @@ func (h *InteractiveHandler) Proxy(ctx context.Context) {
// SessionID: uuid.NewV4().String(),
// IP: asset.Ip,
// port: asset.port,
// UserName: systemUser.UserName,
// Username: systemUser.Username,
// password: systemUser.password,
// PublicKey: parsePrivateKey(systemUser.privateKey)}
//
......@@ -361,10 +363,10 @@ func (h *InteractiveHandler) Proxy(ctx context.Context) {
// nodeConn.Close()
// data := map[string]interface{}{
// "id": nodeConn.SessionID,
// "user": i.user.UserName,
// "user": i.user.Username,
// "asset": asset.Hostname,
// "org_id": asset.OrgID,
// "system_user": systemUser.UserName,
// "system_user": systemUser.Username,
// "login_from": "ST",
// "remote_addr": i.sess.RemoteAddr().String(),
// "is_finished": true,
......@@ -377,10 +379,10 @@ func (h *InteractiveHandler) Proxy(ctx context.Context) {
// }()
// data := map[string]interface{}{
// "id": nodeConn.SessionID,
// "user": i.user.UserName,
// "user": i.user.Username,
// "asset": asset.Hostname,
// "org_id": asset.OrgID,
// "system_user": systemUser.UserName,
// "system_user": systemUser.Username,
// "login_from": "ST",
// "remote_addr": i.sess.RemoteAddr().String(),
// "is_finished": false,
......
package handler
import (
"fmt"
"strings"
)
const (
ColorStart = "\033["
Green = "32m"
Red = "41m"
White = "47m"
ColorEnd = "0m"
Tab = "\t"
EndLine = "\r\n"
Bold = "1"
)
func WrapperString(text string, color string, bold bool) string {
wrapWith := make([]string, 1)
if bold {
wrapWith = append(wrapWith, Bold)
}
wrapWith = append(wrapWith, color)
return fmt.Sprintf("%s%s%s", strings.Join(wrapWith, ";"), text, ColorEnd)
}
func STitle(text string) string {
return WrapperString(text, Green, true)
}
func SWarn(text string) string {
return WrapperString(text, White, false)
}
......@@ -111,7 +111,7 @@ const LoginModeManual = "manual"
type SystemUser struct {
Id string `json:"id"`
Name string `json:"name"`
UserName string `json:"username"`
Username string `json:"username"`
Priority int `json:"priority"`
Protocol string `json:"protocol"`
Comment string `json:"comment"`
......
package proxy
import (
"cocogo/pkg/utils"
"fmt"
"io"
"strconv"
"strings"
"time"
"github.com/gliderlabs/ssh"
"cocogo/pkg/config"
"cocogo/pkg/i18n"
"cocogo/pkg/logger"
......@@ -17,7 +15,7 @@ import (
)
type ProxyServer struct {
Session ssh.Session
UserConn UserConnection
User *model.User
Asset *model.Asset
SystemUser *model.SystemUser
......@@ -33,43 +31,60 @@ func (p *ProxyServer) getSystemUserAuthOrManualSet() {
p.SystemUser.PrivateKey = info.PrivateKey
}
func (p *ProxyServer) checkProtocol() bool {
return true
func (p *ProxyServer) getSystemUserUsernameIfNeed() {
}
func (p *ProxyServer) getSystemUserUsernameIfNeed() {
func (p *ProxyServer) checkProtocolMatch() bool {
return p.SystemUser.Protocol == p.Asset.Protocol
}
func (p *ProxyServer) checkProtocolIsGraph() bool {
switch p.Asset.Protocol {
case "ssh", "telnet":
return true
default:
return false
}
}
func (p *ProxyServer) validatePermission() bool {
return true
}
func (p *ProxyServer) getServerConn() (srvConn ServerConnection, err error) {
func (p *ProxyServer) getSSHConn() (srvConn *ServerSSHConnection, err error) {
srvConn = &ServerSSHConnection{
host: p.Asset.Ip,
port: strconv.Itoa(p.Asset.Port),
user: p.SystemUser.UserName,
user: p.SystemUser.Username,
password: p.SystemUser.Password,
timeout: config.Conf.SSHTimeout,
}
pty, _, ok := p.Session.Pty()
if !ok {
logger.Error("User not request Pty")
return
}
pty := p.UserConn.Pty()
done := make(chan struct{})
go p.sendConnectingMsg(done)
err = srvConn.Connect(pty.Window.Height, pty.Window.Width, pty.Term)
_, _ = io.WriteString(p.Session, "\r\n")
utils.IgnoreErrWriteString(p.UserConn, "\r\n")
close(done)
return
}
func (p *ProxyServer) getTelnetConn() (srvConn *ServerSSHConnection, err error) {
return
}
func (p *ProxyServer) getServerConn() (srvConn ServerConnection, err error) {
if p.Asset.Protocol == "telnet" {
return p.getTelnetConn()
} else {
return p.getSSHConn()
}
}
func (p *ProxyServer) sendConnectingMsg(done chan struct{}) {
delay := 0.0
msg := fmt.Sprintf(i18n.T("Connecting to %s@%s %.1f"), p.SystemUser.UserName, p.Asset.Ip, delay)
_, _ = io.WriteString(p.Session, msg)
msg := fmt.Sprintf(i18n.T("Connecting to %s@%s %.1f"), p.SystemUser.Username, p.Asset.Ip, delay)
utils.IgnoreErrWriteString(p.UserConn, msg)
for int(delay) < config.Conf.SSHTimeout {
select {
case <-done:
......@@ -77,25 +92,47 @@ func (p *ProxyServer) sendConnectingMsg(done chan struct{}) {
default:
delayS := fmt.Sprintf("%.1f", delay)
data := strings.Repeat("\x08", len(delayS)) + delayS
_, _ = io.WriteString(p.Session, data)
utils.IgnoreErrWriteString(p.UserConn, data)
time.Sleep(100 * time.Millisecond)
delay += 0.1
}
}
}
func (p *ProxyServer) preCheckRequisite() (ok bool) {
if !p.checkProtocolMatch() {
msg := utils.WrapperWarn(i18n.T("System user <%s> and asset <%s> protocol are inconsistent."))
msg = fmt.Sprintf(msg, p.SystemUser.Username, p.Asset.Hostname)
utils.IgnoreErrWriteString(p.UserConn, msg)
return
}
if !p.checkProtocolIsGraph() {
msg := i18n.T("Terminal only support protocol ssh/telnet, please use web terminal to access")
msg = utils.WrapperWarn(msg)
utils.IgnoreErrWriteString(p.UserConn, msg)
return
}
if !p.validatePermission() {
msg := fmt.Sprintf("You don't have permission login %s@%s", p.SystemUser.Username, p.Asset.Hostname)
utils.IgnoreErrWriteString(p.UserConn, msg)
return
}
return true
}
func (p *ProxyServer) Proxy() {
if !p.checkProtocol() {
if !p.preCheckRequisite() {
return
}
srvConn, err := p.getServerConn()
if err != nil {
logger.Errorf("Connect host error: %s\n", err)
msg := fmt.Sprintf("Connect asset %s error: %s\n", p.Asset.Hostname, err)
utils.IgnoreErrWriteString(p.UserConn, msg)
logger.Errorf(msg)
return
}
userConn := &UserSSHConnection{Session: p.Session, winch: make(chan ssh.Window)}
sw := NewSwitch(userConn, srvConn)
sw := NewSwitch(p.UserConn, srvConn)
cmdRules, err := service.GetSystemUserFilterRules(p.SystemUser.Id)
if err != nil {
logger.Error("Get system user filter rule error: ", err)
......
......@@ -14,11 +14,11 @@ type UserConnection interface {
User() string
LoginFrom() string
RemoteAddr() string
Pty() ssh.Pty
}
type UserSSHConnection struct {
ssh.Session
winch <-chan ssh.Window
}
func (uc *UserSSHConnection) Protocol() string {
......@@ -30,7 +30,7 @@ func (uc *UserSSHConnection) User() string {
}
func (uc *UserSSHConnection) WinCh() (winch <-chan ssh.Window) {
_, winch, ok := uc.Pty()
_, winch, ok := uc.Session.Pty()
if ok {
return
}
......@@ -44,3 +44,8 @@ func (uc *UserSSHConnection) LoginFrom() string {
func (uc *UserSSHConnection) RemoteAddr() string {
return strings.Split(uc.Session.RemoteAddr().String(), ":")[0]
}
func (uc *UserSSHConnection) Pty() ssh.Pty {
pty, _, _ := uc.Session.Pty()
return pty
}
package utils
import (
"fmt"
"io"
"strings"
)
func IgnoreErrWriteString(writer io.Writer, s string) {
_, _ = io.WriteString(writer, s)
}
const (
ColorEscape = "\033["
Green = "32m"
Red = "31m"
ColorEnd = ColorEscape + "0m"
Bold = "1"
)
const (
CharClear = "\x1b[H\x1b[2J"
CharTab = "\t"
CharNewLine = "\r\n"
)
func WrapperString(text string, color string, meta ...bool) string {
wrapWith := make([]string, 0)
metaLen := len(meta)
switch metaLen {
case 1:
wrapWith = append(wrapWith, Bold)
}
wrapWith = append(wrapWith, color)
return fmt.Sprintf("%s%s%s%s", ColorEscape, strings.Join(wrapWith, ";"), text, ColorEnd)
}
func WrapperTitle(text string) string {
return WrapperString(text, Green, true)
}
func WrapperWarn(text string) string {
text += "\r\n\r\n"
return WrapperString(text, Red)
}
package utils
import (
"fmt"
"testing"
)
func TestWrapperString(t *testing.T) {
s := "Hello world"
s1 := WrapperString(s, Red, true)
fmt.Println(s1)
s2 := WrapperString(s, Green)
fmt.Println(s2)
s3 := WrapperTitle(s)
fmt.Println(s3)
s4 := WrapperWarn(s)
fmt.Println(s4)
}
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