157 |
158 |
175 |
176 |
177 | `
178 |
--------------------------------------------------------------------------------
/gateway/auth_code.go:
--------------------------------------------------------------------------------
1 | /*
2 | * @Copyright Reserved By Janusec (https://www.janusec.com/).
3 | * @Author: U2
4 | * @Date: 2020-05-17 21:45:58
5 | * @Last Modified: U2, 2020-05-17 21:45:58
6 | */
7 |
8 | package gateway
9 |
10 | import (
11 | "encoding/base64"
12 | "fmt"
13 | "net/http"
14 | "strconv"
15 | "text/template"
16 |
17 | "janusec/usermgmt"
18 | "janusec/utils"
19 |
20 | qrcode "github.com/skip2/go-qrcode"
21 | )
22 |
23 | var (
24 | authCodeUITemplate = template.Must(template.New("authcode").Parse(authcodeTemplate))
25 | )
26 |
27 | // AuthCodeContext authenticator code context
28 | type AuthCodeContext struct {
29 | UID string
30 | TOTPKey string
31 | ImageData string
32 | }
33 |
34 | // AuthCodeVerifyFunc Register TOTP in Mobile APP
35 | func AuthCodeVerifyFunc(w http.ResponseWriter, r *http.Request) {
36 | uid := r.FormValue("uid")
37 | totpCode := r.FormValue("code")
38 | totpCodeInt, _ := strconv.ParseUint(totpCode, 10, 32)
39 | totpItem, _ := usermgmt.GetTOTPByUID(uid) //data.DAL.GetTOTPItemByUID(uid)
40 | verifyOK := usermgmt.VerifyCode(totpItem.TOTPKey, uint32(totpCodeInt))
41 | if verifyOK {
42 | _, err := usermgmt.UpdateTOTPVerified(totpItem.ID)
43 | if err != nil {
44 | utils.DebugPrintln("UpdateTOTPVerified error", err)
45 | }
46 | http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
47 | return
48 | }
49 | http.Redirect(w, r, "/oauth/code/register?uid="+uid, http.StatusTemporaryRedirect)
50 | }
51 |
52 | // ShowAuthCodeRegisterUI used for Authenticator Code register UI
53 | func ShowAuthCodeRegisterUI(w http.ResponseWriter, r *http.Request) {
54 | uid := r.FormValue("uid")
55 | totpItem, _ := usermgmt.GetTOTPByUID(uid)
56 | // Format: otpauth://totp/uid?secret=XBSWY3DPEHPK3PXP&issuer=JANUSEC
57 | totpLink := fmt.Sprintf("otpauth://totp/%s?secret=%s&issuer=JANUSEC", uid, totpItem.TOTPKey)
58 | var png []byte
59 | png, err := qrcode.Encode(totpLink, qrcode.Medium, 256)
60 | if err != nil {
61 | utils.DebugPrintln("qrcode.Encode", err)
62 | }
63 | codeImageText := "data:image/png;base64," + base64.StdEncoding.EncodeToString(png)
64 | authCodeContext := AuthCodeContext{
65 | UID: uid,
66 | TOTPKey: totpItem.TOTPKey,
67 | ImageData: codeImageText,
68 | }
69 | if err := authCodeUITemplate.Execute(w, &authCodeContext); err != nil {
70 | http.Error(w, err.Error(), http.StatusInternalServerError)
71 | }
72 | }
73 |
74 | const authcodeTemplate = `
75 |
76 |
77 |
78 |
Authenticator Registration
79 |
80 |
145 |
146 |
147 |
148 |
149 |
150 |
151 |
中文
152 |
English
153 |
154 |
155 |
156 | - Google Authenticator
157 | - Microsoft Authenticator
158 |
159 |
160 |
161 |
162 |
163 |
{{ .TOTPKey }}
164 |
165 |
166 |
167 |
172 |
173 |
176 |
177 | `
178 |
--------------------------------------------------------------------------------
/data/data.go:
--------------------------------------------------------------------------------
1 | /*
2 | * @Copyright Reserved By Janusec (https://www.janusec.com/).
3 | * @Author: U2
4 | * @Date: 2018-07-14 16:23:50
5 | * @Last Modified: U2, 2018-07-14 16:23:50
6 | */
7 |
8 | package data
9 |
10 | import (
11 | "database/sql"
12 | "fmt"
13 | "os"
14 | "strings"
15 |
16 | "janusec/models"
17 | "janusec/utils"
18 |
19 | // PostgreSQL
20 | _ "github.com/lib/pq"
21 | //_ "github.com/mattn/go-sqlite3"
22 | _ "github.com/glebarez/go-sqlite"
23 | )
24 |
25 | // MyDAL used for data access layer
26 | type MyDAL struct {
27 | db *sql.DB
28 | }
29 |
30 | var (
31 | // DAL is Data Access Layer
32 | DAL *MyDAL
33 | // CFG is config
34 | CFG *models.Config
35 | // IsPrimary i.e. Is Primary Node
36 | IsPrimary bool
37 | // Version of JANUSEC
38 | Version = "1.6.0"
39 | )
40 |
41 | // InitConfig init Data Access Layer
42 | func InitConfig() {
43 | DAL = &MyDAL{}
44 | var err error
45 | CFG, err = NewConfig("./config.json")
46 | if err != nil {
47 | utils.DebugPrintln("InitConfig", err)
48 | os.Exit(1)
49 | }
50 | nodeRole := strings.ToLower(CFG.NodeRole)
51 | if nodeRole != "primary" && nodeRole != "replica" {
52 | fmt.Printf("Error: node_role %s is not supported, it should be primary or replica, please check config.json \n", nodeRole)
53 | utils.DebugPrintln("Error: node_role ", nodeRole, " is not supported, it should be primary or replica, please check config.json")
54 | os.Exit(1)
55 | }
56 | IsPrimary = (nodeRole == "primary")
57 | if IsPrimary {
58 | switch CFG.PrimaryNode.DatabaseType {
59 | case "sqlite":
60 | // SQLite
61 | conn := "file:./data.sqlite3?_busy_timeout=9999999"
62 | DAL.db, err = sql.Open("sqlite", conn)
63 | if err != nil {
64 | utils.DebugPrintln("InitConfig sql.Open:", err)
65 | os.Exit(1)
66 | }
67 | // Set max conns 1
68 | DAL.db.SetMaxOpenConns(1)
69 | default:
70 | // PostgreSQL
71 | conn := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=disable",
72 | CFG.PrimaryNode.Database.Host,
73 | CFG.PrimaryNode.Database.Port,
74 | CFG.PrimaryNode.Database.User,
75 | CFG.PrimaryNode.Database.Password,
76 | CFG.PrimaryNode.Database.DBName)
77 | DAL.db, err = sql.Open("postgres", conn)
78 | if err != nil {
79 | utils.DebugPrintln("InitConfig sql.Open:", err)
80 | os.Exit(1)
81 | }
82 | // Check if the User and Password are Correct
83 | _, err = DAL.db.Query("select 1")
84 | if err != nil {
85 | utils.DebugPrintln("InitConfig Failed, Please check the database user and password. Error:", err)
86 | os.Exit(1)
87 | }
88 | // Set max conns 99
89 | DAL.db.SetMaxOpenConns(99)
90 | }
91 | } else {
92 | // Init Nodes Key for replica node
93 | NodesKey = NodeHexKeyToCryptKey(CFG.ReplicaNode.NodeKey)
94 | }
95 | }
96 |
97 | // ExecSQL Exec SQL Directly
98 | func (dal *MyDAL) ExecSQL(sql string) error {
99 | _, err := dal.db.Exec(sql)
100 | return err
101 | }
102 |
103 | // ExistColumnInTable ...
104 | func (dal *MyDAL) ExistColumnInTable(tableName string, columnName string) bool {
105 | var count int64
106 | var sql string
107 | var err error
108 | switch CFG.PrimaryNode.DatabaseType {
109 | case "sqlite":
110 | // SQLite
111 | // This statement has a bug, if a_uid exists, then uid will be considered as exists
112 | // sql = `SELECT count(1) FROM sqlite_master WHERE name=? and sql like ?`
113 | // err = dal.db.QueryRow(sql, tableName, "%"+columnName+"%").Scan(&count)
114 | sql = fmt.Sprintf(`SELECT count(%s) FROM %s`, columnName, tableName)
115 | err = dal.db.QueryRow(sql).Scan(&count)
116 | if err != nil {
117 | // Not exists
118 | return false
119 | }
120 | return true
121 | default:
122 | // PostgreSQL
123 | sql = `SELECT count(1) FROM information_schema.columns WHERE table_name=$1 AND column_name=$2`
124 | err = dal.db.QueryRow(sql, tableName, columnName).Scan(&count)
125 | if err != nil {
126 | utils.DebugPrintln("PostgreSQL ExistColumnInTable QueryRow", tableName, columnName, err)
127 | }
128 | return count > 0
129 | }
130 | }
131 |
132 | // ExistConstraint ...
133 | func (dal *MyDAL) ExistConstraint(tableName string, constraintName string) bool {
134 | var count int64
135 | var sql string
136 | var err error
137 | switch CFG.PrimaryNode.DatabaseType {
138 | case "sqlite":
139 | // SQLite
140 | // select * from sqlite_master where type='index' and tbl_name='test' and name='uid'
141 | // For SQLite, create unique index uid on table_name(column1, column2);
142 | sql = `SELECT count(1) FROM sqlite_master WHERE type='index' AND tbl_name=$1 AND name=$2`
143 | err = dal.db.QueryRow(sql, tableName, constraintName).Scan(&count)
144 | if err != nil {
145 | utils.DebugPrintln("ExistConstraint QueryRow", err)
146 | }
147 | return count > 0
148 | default:
149 | // PostgreSQL
150 | sql = `SELECT count(1) FROM information_schema.constraint_column_usage WHERE table_name=$1 and constraint_name=$2`
151 | err = dal.db.QueryRow(sql, tableName, constraintName).Scan(&count)
152 | if err != nil {
153 | utils.DebugPrintln("ExistConstraint QueryRow", err)
154 | }
155 | return count > 0
156 | }
157 | }
158 |
--------------------------------------------------------------------------------
/data/firewall_group_policy.go:
--------------------------------------------------------------------------------
1 | /*
2 | * @Copyright Reserved By Janusec (https://www.janusec.com/).
3 | * @Author: U2
4 | * @Date: 2018-07-14 16:31:06
5 | * @Last Modified: U2, 2018-07-14 16:31:06
6 | */
7 |
8 | package data
9 |
10 | import (
11 | "janusec/models"
12 | "janusec/utils"
13 | )
14 |
15 | const (
16 | sqlCreateTableIfNotExistsGroupPolicy = `CREATE TABLE IF NOT EXISTS "group_policies"("id" bigserial primary key,"description" VARCHAR(256) NOT NULL DEFAULT '',"app_id" bigint,"vuln_id" bigint,"hit_value" bigint,"action" bigint,"is_enabled" boolean,"user_id" bigint,"update_time" bigint)`
17 | sqlExistsGroupPolicy = `SELECT COALESCE((SELECT 1 FROM "group_policies" limit 1),0)`
18 | sqlSelectGroupPolicies = `SELECT "id","description","app_id","vuln_id","hit_value","action","is_enabled","user_id","update_time" FROM "group_policies"`
19 | sqlSelectGroupPoliciesByAppID = `SELECT "id","description","vuln_id","hit_value","action","is_enabled","user_id","update_time" FROM "group_policies" WHERE "app_id"=$1`
20 | sqlInsertGroupPolicy = `INSERT INTO "group_policies"("id","description","app_id","vuln_id","hit_value","action","is_enabled","user_id","update_time") VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9) RETURNING "id"`
21 | sqlUpdateGroupPolicy = `UPDATE "group_policies" SET "description"=$1,"app_id"=$2,"vuln_id"=$3,"hit_value"=$4,"action"=$5,"is_enabled"=$6,"user_id"=$7,"update_time"=$8 WHERE "id"=$9`
22 | sqlDeleteGroupPolicyByID = `DELETE FROM "group_policies" WHERE "id"=$1`
23 | )
24 |
25 | // CreateTableIfNotExistsGroupPolicy ...
26 | func (dal *MyDAL) CreateTableIfNotExistsGroupPolicy() error {
27 | _, err := dal.db.Exec(sqlCreateTableIfNotExistsGroupPolicy)
28 | if err != nil {
29 | utils.DebugPrintln("CreateTableIfNotExistsGroupPolicy", err)
30 | }
31 | return err
32 | }
33 |
34 | // DeleteGroupPolicyByID ...
35 | func (dal *MyDAL) DeleteGroupPolicyByID(id int64) error {
36 | _, err := dal.db.Exec(sqlDeleteGroupPolicyByID, id)
37 | if err != nil {
38 | utils.DebugPrintln("DeleteGroupPolicyByID", err)
39 | }
40 | return err
41 | }
42 |
43 | // UpdateGroupPolicy ...
44 | func (dal *MyDAL) UpdateGroupPolicy(description string, appID int64, vulnID int64, hitValue int64, action models.PolicyAction, isEnabled bool, userID int64, updateTime int64, id int64) error {
45 | stmt, _ := dal.db.Prepare(sqlUpdateGroupPolicy)
46 | defer stmt.Close()
47 | _, err := stmt.Exec(description, appID, vulnID, hitValue, action, isEnabled, userID, updateTime, id)
48 | if err != nil {
49 | utils.DebugPrintln("UpdateGroupPolicy", err)
50 | }
51 | return err
52 | }
53 |
54 | // SelectGroupPolicies ...
55 | func (dal *MyDAL) SelectGroupPolicies() []*models.GroupPolicy {
56 | groupPolicies := []*models.GroupPolicy{}
57 | rows, err := dal.db.Query(sqlSelectGroupPolicies)
58 | if err != nil {
59 | utils.DebugPrintln("SelectGroupPolicies", err)
60 | }
61 | defer rows.Close()
62 | for rows.Next() {
63 | groupPolicy := &models.GroupPolicy{}
64 | err = rows.Scan(&groupPolicy.ID, &groupPolicy.Description, &groupPolicy.AppID, &groupPolicy.VulnID,
65 | &groupPolicy.HitValue, &groupPolicy.Action, &groupPolicy.IsEnabled, &groupPolicy.UserID, &groupPolicy.UpdateTime)
66 | if err != nil {
67 | utils.DebugPrintln("SelectGroupPolicies Scan", err)
68 | }
69 | groupPolicies = append(groupPolicies, groupPolicy)
70 | }
71 | return groupPolicies
72 | }
73 |
74 | // SelectGroupPoliciesByAppID ...
75 | func (dal *MyDAL) SelectGroupPoliciesByAppID(appID int64) ([]*models.GroupPolicy, error) {
76 | groupPolicies := []*models.GroupPolicy{}
77 | rows, err := dal.db.Query(sqlSelectGroupPoliciesByAppID, appID)
78 | if err != nil {
79 | utils.DebugPrintln("SelectGroupPoliciesByAppID", err)
80 | }
81 | defer rows.Close()
82 | for rows.Next() {
83 | groupPolicy := &models.GroupPolicy{}
84 | groupPolicy.AppID = appID
85 | err = rows.Scan(&groupPolicy.ID, &groupPolicy.Description, &groupPolicy.VulnID,
86 | &groupPolicy.HitValue, &groupPolicy.Action, &groupPolicy.IsEnabled, &groupPolicy.UserID, &groupPolicy.UpdateTime)
87 | if err != nil {
88 | utils.DebugPrintln("SelectGroupPoliciesByAppID Scan", err)
89 | return groupPolicies, err
90 | }
91 | groupPolicies = append(groupPolicies, groupPolicy)
92 | }
93 | return groupPolicies, err
94 | }
95 |
96 | // InsertGroupPolicy ...
97 | func (dal *MyDAL) InsertGroupPolicy(description string, appID int64, vulnID int64, hitValue int64, action models.PolicyAction, isEnabled bool, userID int64, updateTime int64) (newID int64, err error) {
98 | stmt, err := dal.db.Prepare(sqlInsertGroupPolicy)
99 | if err != nil {
100 | utils.DebugPrintln("InsertGroupPolicy Prepare", err)
101 | }
102 | defer stmt.Close()
103 | id := utils.GenSnowflakeID()
104 | err = stmt.QueryRow(id, description, appID, vulnID, hitValue, action, isEnabled, userID, updateTime).Scan(&newID)
105 | if err != nil {
106 | utils.DebugPrintln("InsertGroupPolicy Scan", err)
107 | }
108 | return newID, err
109 | }
110 |
111 | // ExistsGroupPolicy ...
112 | func (dal *MyDAL) ExistsGroupPolicy() bool {
113 | var exist int
114 | err := dal.db.QueryRow(sqlExistsGroupPolicy).Scan(&exist)
115 | if err != nil {
116 | utils.DebugPrintln("ExistsGroupPolicy", err)
117 | }
118 | return exist != 0
119 | }
120 |
--------------------------------------------------------------------------------
/gateway/auth_oauth.go:
--------------------------------------------------------------------------------
1 | /*
2 | * @Copyright Reserved By Janusec (https://www.janusec.com/).
3 | * @Author: U2
4 | * @Date: 2020-03-14 18:47:18
5 | * @Last Modified: U2, 2020-03-14 18:47:18
6 | */
7 |
8 | package gateway
9 |
10 | import (
11 | "fmt"
12 | "net/http"
13 |
14 | "janusec/data"
15 | "janusec/usermgmt"
16 | )
17 |
18 | // OAuthInfo OAuth Information
19 | type OAuthInfo struct {
20 | UseOAuth bool `json:"use_oauth"`
21 | DisplayName string `json:"display_name"`
22 | EntranceURL string `json:"entrance_url"`
23 |
24 | // AuthenticatorEnabled added in v1.2.2, for janusec-admin login
25 | AuthenticatorEnabled bool `json:"authenticator_enabled"`
26 | }
27 |
28 | // WxworkCallBackHandleFunc for Wxwork CallBack
29 | func WxworkCallBackHandleFunc(w http.ResponseWriter, r *http.Request) {
30 | usermgmt.WxworkCallbackWithCode(w, r)
31 | }
32 |
33 | // DingtalkCallBackHandleFunc for Dingtalk CallBack
34 | func DingtalkCallBackHandleFunc(w http.ResponseWriter, r *http.Request) {
35 | usermgmt.DingtalkCallbackWithCode(w, r)
36 | }
37 |
38 | // FeishuCallBackHandleFunc for Feishu CallBack
39 | func FeishuCallBackHandleFunc(w http.ResponseWriter, r *http.Request) {
40 | usermgmt.FeishuCallbackWithCode(w, r)
41 | }
42 |
43 | // LarkCallBackHandleFunc for Lark CallBack
44 | func LarkCallBackHandleFunc(w http.ResponseWriter, r *http.Request) {
45 | usermgmt.LarkCallbackWithCode(w, r)
46 | }
47 |
48 | // CAS2CallBackHandleFunc for CAS2 CallBack
49 | func CAS2CallBackHandleFunc(w http.ResponseWriter, r *http.Request) {
50 | usermgmt.CAS2CallbackWithCode(w, r)
51 | }
52 |
53 | // LDAPCallBackHandleFunc for LDAP CallBack
54 | func LDAPCallBackHandleFunc(w http.ResponseWriter, r *http.Request) {
55 | usermgmt.LDAPAuthFunc(w, r)
56 | }
57 |
58 | // OAuthGetHandleFunc Get OAuth Information and Response
59 | func OAuthGetHandleFunc(w http.ResponseWriter, r *http.Request) {
60 | obj, err := GetOAuthInfo()
61 | GenResponseByObject(w, obj, err)
62 | }
63 |
64 | // GetOAuthInfo Get OAuth Information
65 | func GetOAuthInfo() (*OAuthInfo, error) {
66 | oauthInfo := OAuthInfo{}
67 | oauthInfo.AuthenticatorEnabled = data.PrimarySetting.AuthenticatorEnabled
68 | if !data.NodeSetting.AuthConfig.Enabled {
69 | return &oauthInfo, nil
70 | }
71 | switch data.NodeSetting.AuthConfig.Provider {
72 | case "wxwork":
73 | entranceURL := fmt.Sprintf("https://open.work.weixin.qq.com/wwopen/sso/qrConnect?appid=%s&agentid=%s&redirect_uri=%s&state=admin",
74 | data.NodeSetting.AuthConfig.Wxwork.CorpID,
75 | data.NodeSetting.AuthConfig.Wxwork.AgentID,
76 | data.NodeSetting.AuthConfig.Wxwork.Callback)
77 | oauthInfo.UseOAuth = true
78 | oauthInfo.DisplayName = data.NodeSetting.AuthConfig.Wxwork.DisplayName
79 | oauthInfo.EntranceURL = entranceURL
80 | return &oauthInfo, nil
81 | case "dingtalk":
82 | /* API v1
83 | entranceURL := fmt.Sprintf("https://oapi.dingtalk.com/connect/qrconnect?appid=%s&response_type=code&scope=snsapi_login&state=admin&redirect_uri=%s",
84 | data.NodeSetting.AuthConfig.Dingtalk.AppID,
85 | data.NodeSetting.AuthConfig.Dingtalk.Callback)
86 | */
87 | // API v2, added on Mar 23, 2024, v1.5.0
88 | entranceURL := fmt.Sprintf(`https://login.dingtalk.com/oauth2/auth?redirect_uri=%s&response_type=code&corpId=%s&client_id=%s&scope=openid%%20corpid&state=admin&prompt=consent`,
89 | data.NodeSetting.AuthConfig.Dingtalk.Callback,
90 | data.NodeSetting.AuthConfig.Dingtalk.CorpID,
91 | data.NodeSetting.AuthConfig.Dingtalk.AppID)
92 | oauthInfo.UseOAuth = true
93 | oauthInfo.DisplayName = data.NodeSetting.AuthConfig.Dingtalk.DisplayName
94 | oauthInfo.EntranceURL = entranceURL
95 | return &oauthInfo, nil
96 | case "feishu":
97 | entranceURL := fmt.Sprintf("https://open.feishu.cn/open-apis/authen/v1/index?redirect_uri=%s&app_id=%s&state=admin",
98 | data.NodeSetting.AuthConfig.Feishu.Callback,
99 | data.NodeSetting.AuthConfig.Feishu.AppID)
100 | oauthInfo.UseOAuth = true
101 | oauthInfo.DisplayName = data.NodeSetting.AuthConfig.Feishu.DisplayName
102 | oauthInfo.EntranceURL = entranceURL
103 | return &oauthInfo, nil
104 | case "lark":
105 | entranceURL := fmt.Sprintf("https://open.larksuite.com/open-apis/authen/v1/index?redirect_uri=%s&app_id=%s&state=admin",
106 | data.NodeSetting.AuthConfig.Lark.Callback,
107 | data.NodeSetting.AuthConfig.Lark.AppID)
108 | oauthInfo.UseOAuth = true
109 | oauthInfo.DisplayName = data.NodeSetting.AuthConfig.Lark.DisplayName
110 | oauthInfo.EntranceURL = entranceURL
111 | return &oauthInfo, nil
112 | case "ldap":
113 | entranceURL := data.NodeSetting.AuthConfig.LDAP.Entrance + "?state=admin"
114 | oauthInfo.UseOAuth = true
115 | oauthInfo.DisplayName = data.NodeSetting.AuthConfig.LDAP.DisplayName
116 | oauthInfo.EntranceURL = entranceURL
117 | return &oauthInfo, nil
118 | case "cas2":
119 | entranceURL := fmt.Sprintf("%s/login?renew=true&service=%s?state=admin",
120 | data.NodeSetting.AuthConfig.CAS2.Entrance, data.NodeSetting.AuthConfig.CAS2.Callback)
121 | oauthInfo.UseOAuth = true
122 | oauthInfo.DisplayName = data.NodeSetting.AuthConfig.CAS2.DisplayName
123 | oauthInfo.EntranceURL = entranceURL
124 | return &oauthInfo, nil
125 | }
126 | oauthInfo.UseOAuth = false
127 | return &oauthInfo, nil // errors.New("No OAuth2 provider, you can enable it in config.json")
128 | }
129 |
--------------------------------------------------------------------------------
/usermgmt/oauth_feishu.go:
--------------------------------------------------------------------------------
1 | /*
2 | * @Copyright Reserved By Janusec (https://www.janusec.com/).
3 | * @Author: U2
4 | * @Date: 2020-03-23 21:02:39
5 | * @Last Modified: U2, 2020-03-23 21:02:39
6 | */
7 |
8 | package usermgmt
9 |
10 | import (
11 | "bytes"
12 | "encoding/json"
13 | "fmt"
14 | "net/http"
15 |
16 | "janusec/utils"
17 |
18 | "janusec/data"
19 | "janusec/models"
20 |
21 | "github.com/gorilla/sessions"
22 | "github.com/patrickmn/go-cache"
23 | )
24 |
25 | type FeishuAccessToken struct {
26 | Code int64 `json:"code"`
27 | Msg string `json:"msg"`
28 | AppAccessToken string `json:"app_access_token"`
29 | Expire int `json:"expire"`
30 | }
31 |
32 | // https://open.feishu.cn/document/ukTMukTMukTM/uEDO4UjLxgDO14SM4gTN
33 | type FeishuUserReqBody struct {
34 | AppAccessToken string `json:"app_access_token"`
35 | GrantType string `json:"grant_type"`
36 | Code string `json:"code"`
37 | }
38 |
39 | // https://open.feishu.cn/document/ukTMukTMukTM/uEDO4UjLxgDO14SM4gTN
40 | type FeishuUser struct {
41 | Code int64 `json:"code"`
42 | Msg string `json:"msg"`
43 | Data FeishuAuthData `json:"data"`
44 | }
45 |
46 | type FeishuAuthData struct {
47 | AccessToken string `json:"access_token"`
48 | EnName string `json:"en_name"`
49 | }
50 |
51 | // Doc: https://open.feishu.cn/document/ukTMukTMukTM/ukzN4UjL5cDO14SO3gTN
52 | // Step 1: GET https://open.feishu.cn/open-apis/authen/v1/index?redirect_uri={REDIRECT_URI}&app_id={APPID}&state={STATE}
53 | // If state==admin, for janusec-admin; else for frontend applications
54 | func FeishuCallbackWithCode(w http.ResponseWriter, r *http.Request) {
55 | // Step 2.1: Callback with code and state, http://gate.janusec.com/?code=BM8k8U6RwtQtNY&state=admin
56 | code := r.FormValue("code")
57 | state := r.FormValue("state")
58 | // Step 2.2: Within Callback, get app_access_token
59 | // Doc: https://open.feishu.cn/document/ukTMukTMukTM/uADN14CM0UjLwQTN
60 | // POST https://open.feishu.cn/open-apis/auth/v3/app_access_token/internal/
61 | // {"app_id":"cli_slkdasd", "app_secret":"dskLLdkasdKK"}
62 | // accessTokenURL := "https://open.feishu.cn/open-apis/auth/v3/app_access_token/internal/"
63 | body := fmt.Sprintf(`{"app_id":"%s", "app_secret":"%s"}`,
64 | data.NodeSetting.AuthConfig.Feishu.AppID,
65 | data.NodeSetting.AuthConfig.Feishu.AppSecret)
66 | request, _ := http.NewRequest("POST",
67 | "https://open.feishu.cn/open-apis/auth/v3/app_access_token/internal/",
68 | bytes.NewReader([]byte(body)))
69 | resp, err := utils.GetResponse(request)
70 | if err != nil {
71 | utils.DebugPrintln("FeishuCallbackWithCode GetResponse", err)
72 | }
73 | tokenResponse := FeishuAccessToken{}
74 | err = json.Unmarshal(resp, &tokenResponse)
75 | if err != nil {
76 | utils.DebugPrintln("FeishuCallbackWithCode json.Unmarshal error", err)
77 | }
78 | // Step 2.3: Get User name
79 | // https://open.feishu.cn/document/ukTMukTMukTM/uEDO4UjLxgDO14SM4gTN
80 | userURL := "https://open.feishu.cn/open-apis/authen/v1/access_token"
81 | feishuUserReqBody := FeishuUserReqBody{
82 | AppAccessToken: tokenResponse.AppAccessToken,
83 | GrantType: "authorization_code",
84 | Code: code,
85 | }
86 | bytesData, _ := json.Marshal(feishuUserReqBody)
87 | request, _ = http.NewRequest("POST", userURL, bytes.NewReader(bytesData))
88 | request.Header.Set("Content-Type", "application/json")
89 |
90 | resp, err = utils.GetResponse(request)
91 | if err != nil {
92 | utils.DebugPrintln("FeishuCallbackWithCode GetResponse", err)
93 | }
94 | feishuUser := FeishuUser{}
95 | err = json.Unmarshal(resp, &feishuUser)
96 | if err != nil {
97 | utils.DebugPrintln("FeishuCallbackWithCode json.Unmarshal error", err)
98 | }
99 | if state == "admin" {
100 | appUser := data.DAL.SelectAppUserByName(feishuUser.Data.EnName)
101 | var userID int64
102 | if appUser == nil {
103 | // Insert into db if not existed
104 | userID, err = data.DAL.InsertIfNotExistsAppUser(feishuUser.Data.EnName, "", "", "", false, false, false, false)
105 | if err != nil {
106 | w.WriteHeader(403)
107 | w.Write([]byte("Error: " + err.Error()))
108 | return
109 | }
110 | } else {
111 | userID = appUser.ID
112 | }
113 | // create session
114 | authUser := &models.AuthUser{
115 | UserID: userID,
116 | Username: feishuUser.Data.EnName,
117 | Logged: true,
118 | IsSuperAdmin: appUser.IsSuperAdmin,
119 | IsCertAdmin: appUser.IsCertAdmin,
120 | IsAppAdmin: appUser.IsAppAdmin,
121 | NeedModifyPWD: false}
122 | session, _ := store.Get(r, "sessionid")
123 | session.Values["authuser"] = authUser
124 | session.Options = &sessions.Options{Path: "/janusec-admin/", MaxAge: tokenResponse.Expire}
125 | err = session.Save(r, w)
126 | if err != nil {
127 | utils.DebugPrintln("FeishuCallbackWithCode session save error", err)
128 | }
129 | RecordAuthLog(r, authUser.Username, "Feishu", data.CFG.PrimaryNode.Admin.Portal)
130 | http.Redirect(w, r, data.CFG.PrimaryNode.Admin.Portal, http.StatusTemporaryRedirect)
131 | return
132 | }
133 | // Gateway OAuth for employees and internal application
134 | oauthStateI, found := OAuthCache.Get(state)
135 | if found {
136 | oauthState := oauthStateI.(models.OAuthState)
137 | oauthState.UserID = feishuUser.Data.EnName
138 | oauthState.AccessToken = feishuUser.Data.AccessToken
139 | OAuthCache.Set(state, oauthState, cache.DefaultExpiration)
140 | RecordAuthLog(r, oauthState.UserID, "Feishu", oauthState.CallbackURL)
141 | http.Redirect(w, r, oauthState.CallbackURL, http.StatusTemporaryRedirect)
142 | return
143 | }
144 | http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
145 | }
146 |
--------------------------------------------------------------------------------
/data/backend_setting.go:
--------------------------------------------------------------------------------
1 | /*
2 | * @Copyright Reserved By Janusec (https://www.janusec.com/).
3 | * @Author: U2
4 | * @Date: 2018-07-14 16:24:56
5 | * @Last Modified: U2, 2018-07-14 16:24:56
6 | */
7 |
8 | package data
9 |
10 | import (
11 | "janusec/utils"
12 | )
13 |
14 | const (
15 | sqlCreateTableIfNotExistsSettings = `CREATE TABLE IF NOT EXISTS "settings"("id" bigserial PRIMARY KEY, "name" VARCHAR(128) NOT NULL,"bool_value" boolean,"int_value" bigint,"float_value" decimal,"string_value" VARCHAR(16384))`
16 | sqlCountSettings = `SELECT COUNT(1) FROM "settings"`
17 | sqlInsertBoolSetting = `INSERT INTO "settings"("id","name","bool_value") VALUES($1,$2,$3)`
18 | sqlInsertIntSetting = `INSERT INTO "settings"("id","name","int_value") VALUES($1,$2,$3)`
19 | sqlInsertFloatSetting = `INSERT INTO "settings"("id","name","float_value") VALUES($1,$2,$3)`
20 | sqlInsertStringSetting = `INSERT INTO "settings"("id","name","string_value") VALUES($1,$2,$3)`
21 | sqlUpdateBoolSetting = `UPDATE "settings" SET "bool_value"=$1 WHERE "name"=$2`
22 | sqlUpdateIntSetting = `UPDATE "settings" SET "int_value"=$1 WHERE "name"=$2`
23 | sqlUpdateFloatSetting = `UPDATE "settings" SET "float_value"=$1 WHERE "name"=$2`
24 | sqlUpdateStringSetting = `UPDATE "settings" SET "string_value"=$1 WHERE "name"=$2`
25 | sqlSelectBoolSetting = `SELECT "bool_value" FROM "settings" WHERE "name"=$1`
26 | sqlSelectIntSetting = `SELECT "int_value" FROM "settings" WHERE "name"=$1`
27 | sqlSelectFloatSetting = `SELECT "float_value" FROM "settings" WHERE "name"=$1`
28 | sqlSelectStringSetting = `SELECT "string_value" FROM "settings" WHERE "name"=$1`
29 | sqlExistsSetting = `SELECT COALESCE((SELECT 1 FROM "settings" WHERE "name"=$1 limit 1),0)`
30 | )
31 |
32 | // ExistsSetting ...
33 | func (dal *MyDAL) ExistsSetting(name string) bool {
34 | var exist int
35 | err := dal.db.QueryRow(sqlExistsSetting, name).Scan(&exist)
36 | if err != nil {
37 | utils.DebugPrintln("Check ExistsSetting: "+name, err)
38 | }
39 | return exist != 0
40 | }
41 |
42 | // SelectBoolSetting ...
43 | func (dal *MyDAL) SelectBoolSetting(name string) (value bool) {
44 | err := dal.db.QueryRow(sqlSelectBoolSetting, name).Scan(&value)
45 | if err != nil {
46 | utils.DebugPrintln("SelectBoolSetting: "+name, err)
47 | }
48 | return value
49 | }
50 |
51 | // SelectIntSetting ...
52 | func (dal *MyDAL) SelectIntSetting(name string) (value int64) {
53 | err := dal.db.QueryRow(sqlSelectIntSetting, name).Scan(&value)
54 | if err != nil {
55 | utils.DebugPrintln("SelectIntSetting: "+name, err)
56 | }
57 | return value
58 | }
59 |
60 | // SelectFloatSetting ...
61 | func (dal *MyDAL) SelectFloatSetting(name string) (value float64) {
62 | err := dal.db.QueryRow(sqlSelectFloatSetting, name).Scan(&value)
63 | if err != nil {
64 | utils.DebugPrintln("SelectFloatSetting: "+name, err)
65 | }
66 | return value
67 | }
68 |
69 | // SelectStringSetting ...
70 | func (dal *MyDAL) SelectStringSetting(name string) (value string) {
71 | err := dal.db.QueryRow(sqlSelectStringSetting, name).Scan(&value)
72 | if err != nil {
73 | utils.DebugPrintln("SelectStringSetting: "+name, err)
74 | }
75 | return value
76 | }
77 |
78 | // SaveBoolSetting ...
79 | func (dal *MyDAL) SaveBoolSetting(name string, value bool) (err error) {
80 | if dal.ExistsSetting(name) {
81 | _, err = dal.db.Exec(sqlUpdateBoolSetting, value, name)
82 | } else {
83 | id := utils.GenSnowflakeID()
84 | _, err = dal.db.Exec(sqlInsertBoolSetting, id, name, value)
85 | }
86 | if err != nil {
87 | utils.DebugPrintln("SaveBoolSetting: "+name, err)
88 | }
89 | return err
90 | }
91 |
92 | // SaveIntSetting ...
93 | func (dal *MyDAL) SaveIntSetting(name string, value int64) (err error) {
94 | if dal.ExistsSetting(name) {
95 | _, err = dal.db.Exec(sqlUpdateIntSetting, value, name)
96 | } else {
97 | id := utils.GenSnowflakeID()
98 | _, err = dal.db.Exec(sqlInsertIntSetting, id, name, value)
99 | }
100 | if err != nil {
101 | utils.DebugPrintln("SaveIntSetting: "+name, err)
102 | }
103 | return err
104 | }
105 |
106 | // SaveFloatSetting ...
107 | func (dal *MyDAL) SaveFloatSetting(name string, value float64) (err error) {
108 | if dal.ExistsSetting(name) {
109 | _, err = dal.db.Exec(sqlUpdateFloatSetting, value, name)
110 | } else {
111 | id := utils.GenSnowflakeID()
112 | _, err = dal.db.Exec(sqlInsertFloatSetting, id, name, value)
113 | }
114 | if err != nil {
115 | utils.DebugPrintln("SaveFloatSetting: "+name, err)
116 | }
117 | return err
118 | }
119 |
120 | // SaveStringSetting ...
121 | func (dal *MyDAL) SaveStringSetting(name string, value string) (err error) {
122 | if dal.ExistsSetting(name) {
123 | _, err = dal.db.Exec(sqlUpdateStringSetting, value, name)
124 | } else {
125 | id := utils.GenSnowflakeID()
126 | _, err = dal.db.Exec(sqlInsertStringSetting, id, name, value)
127 | }
128 | if err != nil {
129 | utils.DebugPrintln("SaveStringSetting: "+name, err)
130 | }
131 | return err
132 | }
133 |
134 | // CreateTableIfNotExistsSettings ...
135 | func (dal *MyDAL) CreateTableIfNotExistsSettings() error {
136 | _, err := dal.db.Exec(sqlCreateTableIfNotExistsSettings)
137 | if err != nil {
138 | utils.DebugPrintln("CreateTableIfNotExistsSettings", err)
139 | }
140 | return err
141 | }
142 |
143 | // CountSettings ...
144 | func (dal *MyDAL) CountSettings() int64 {
145 | var settingsCount int64
146 | err := dal.db.QueryRow(sqlCountSettings).Scan(&settingsCount)
147 | if err != nil {
148 | utils.DebugPrintln("CountSettings", err)
149 | }
150 | return settingsCount
151 | }
152 |
--------------------------------------------------------------------------------
/usermgmt/oauth_lark.go:
--------------------------------------------------------------------------------
1 | /*
2 | * @Copyright Reserved By Janusec (https://www.janusec.com/).
3 | * @Author: U2
4 | * @Date: 2020-03-23 21:02:39
5 | * @Last Modified: U2, 2020-03-23 21:02:39
6 | */
7 |
8 | package usermgmt
9 |
10 | import (
11 | "bytes"
12 | "encoding/json"
13 | "fmt"
14 | "net/http"
15 |
16 | "janusec/utils"
17 |
18 | "janusec/data"
19 | "janusec/models"
20 |
21 | "github.com/gorilla/sessions"
22 | "github.com/patrickmn/go-cache"
23 | )
24 |
25 | type LarkAccessToken struct {
26 | Code int64 `json:"code"`
27 | Msg string `json:"msg"`
28 | AppAccessToken string `json:"app_access_token"`
29 | Expire int `json:"expire"`
30 | }
31 |
32 | // https://open.larksuite.com/document/ukTMukTMukTM/uEDO4UjLxgDO14SM4gTN
33 | type LarkUserReqBody struct {
34 | AppAccessToken string `json:"app_access_token"`
35 | GrantType string `json:"grant_type"`
36 | Code string `json:"code"`
37 | }
38 |
39 | // https://open.larksuite.com/document/uMzMyEjLzMjMx4yMzITM/ukTN0EjL5UDNx4SO1QTM
40 | type LarkUser struct {
41 | Code int64 `json:"code"`
42 | Msg string `json:"msg"`
43 | Data LarkAuthData `json:"data"`
44 | }
45 |
46 | type LarkAuthData struct {
47 | AccessToken string `json:"access_token"`
48 | EnName string `json:"en_name"`
49 | }
50 |
51 | // Doc: https://open.larksuite.com/document/uMzMyEjLzMjMx4yMzITM/ugTN0EjL4UDNx4CO1QTM
52 | // Step 1: GET https://open.larksuite.com/open-apis/authen/v1/index?redirect_uri={REDIRECT_URI}&app_id={APPID}&state={STATE}
53 | // If state==admin, for janusec-admin; else for frontend applications
54 | func LarkCallbackWithCode(w http.ResponseWriter, r *http.Request) {
55 | // Step 2.1: Callback with code and state, http://gate.janusec.com/?code=BM8k8U6RwtQtNY&state=admin
56 | code := r.FormValue("code")
57 | state := r.FormValue("state")
58 | // Step 2.2: Within Callback, get app_access_token
59 | // Doc: https://open.larksuite.com/document/uMzMyEjLzMjMx4yMzITM/uMjN0EjLzYDNx4yM2QTM
60 | // POST https://open.larksuite.com/open-apis/auth/v3/app_access_token/internal/
61 | // {"app_id":"cli_slkdasd", "app_secret":"dskLLdkasdKK"}
62 | // accessTokenURL := "https://open.larksuite.com/open-apis/auth/v3/app_access_token/internal/"
63 | body := fmt.Sprintf(`{"app_id":"%s", "app_secret":"%s"}`,
64 | data.NodeSetting.AuthConfig.Lark.AppID,
65 | data.NodeSetting.AuthConfig.Lark.AppSecret)
66 | request, _ := http.NewRequest("POST",
67 | "https://open.larksuite.com/open-apis/auth/v3/app_access_token/internal",
68 | bytes.NewReader([]byte(body)))
69 | resp, err := utils.GetResponse(request)
70 | if err != nil {
71 | utils.DebugPrintln("LarkCallbackWithCode GetResponse", err)
72 | }
73 | tokenResponse := LarkAccessToken{}
74 | err = json.Unmarshal(resp, &tokenResponse)
75 | if err != nil {
76 | utils.DebugPrintln("LarkCallbackWithCode json.Unmarshal error", err)
77 | }
78 | // Step 2.3: Get User name
79 | // https://open.larksuite.com/document/uMzMyEjLzMjMx4yMzITM/ukTN0EjL5UDNx4SO1QTM
80 | userURL := "https://open.larksuite.com/open-apis/authen/v1/access_token"
81 | larkUserReqBody := LarkUserReqBody{
82 | AppAccessToken: tokenResponse.AppAccessToken,
83 | GrantType: "authorization_code",
84 | Code: code,
85 | }
86 | bytesData, err := json.Marshal(larkUserReqBody)
87 | if err != nil {
88 | utils.DebugPrintln("LarkCallbackWithCode json.Marshal", err)
89 | }
90 | request, err = http.NewRequest("POST", userURL, bytes.NewReader(bytesData))
91 | if err != nil {
92 | utils.DebugPrintln("LarkCallbackWithCode http.NewRequest", err)
93 | }
94 | request.Header.Set("Content-Type", "application/json")
95 |
96 | resp, err = utils.GetResponse(request)
97 | if err != nil {
98 | utils.DebugPrintln("LarkCallbackWithCode GetResponse", err)
99 | }
100 | larkUser := LarkUser{}
101 | err = json.Unmarshal(resp, &larkUser)
102 | if err != nil {
103 | utils.DebugPrintln("LarkCallbackWithCode json.Unmarshal error", err)
104 | }
105 | if state == "admin" {
106 | appUser := data.DAL.SelectAppUserByName(larkUser.Data.EnName)
107 | var userID int64
108 | if appUser == nil {
109 | // Insert into db if not existed
110 | userID, err = data.DAL.InsertIfNotExistsAppUser(larkUser.Data.EnName, "", "", "", false, false, false, false)
111 | if err != nil {
112 | w.WriteHeader(403)
113 | w.Write([]byte("Error: " + err.Error()))
114 | return
115 | }
116 | } else {
117 | userID = appUser.ID
118 | }
119 | // create session
120 | authUser := &models.AuthUser{
121 | UserID: userID,
122 | Username: larkUser.Data.EnName,
123 | Logged: true,
124 | IsSuperAdmin: appUser.IsSuperAdmin,
125 | IsCertAdmin: appUser.IsCertAdmin,
126 | IsAppAdmin: appUser.IsAppAdmin,
127 | NeedModifyPWD: false}
128 | session, _ := store.Get(r, "sessionid")
129 | session.Values["authuser"] = authUser
130 | session.Options = &sessions.Options{Path: "/janusec-admin/", MaxAge: tokenResponse.Expire}
131 | err = session.Save(r, w)
132 | if err != nil {
133 | utils.DebugPrintln("LarkCallbackWithCode session save error", err)
134 | }
135 | RecordAuthLog(r, authUser.Username, "Lark", data.CFG.PrimaryNode.Admin.Portal)
136 | http.Redirect(w, r, data.CFG.PrimaryNode.Admin.Portal, http.StatusTemporaryRedirect)
137 | return
138 | }
139 | // Gateway OAuth for employees and internal application
140 | oauthStateI, found := OAuthCache.Get(state)
141 | if found {
142 | oauthState := oauthStateI.(models.OAuthState)
143 | oauthState.UserID = larkUser.Data.EnName
144 | oauthState.AccessToken = larkUser.Data.AccessToken
145 | OAuthCache.Set(state, oauthState, cache.DefaultExpiration)
146 | RecordAuthLog(r, oauthState.UserID, "Lark", oauthState.CallbackURL)
147 | http.Redirect(w, r, oauthState.CallbackURL, http.StatusTemporaryRedirect)
148 | return
149 | }
150 | http.Redirect(w, r, "/", http.StatusTemporaryRedirect)
151 | }
152 |
--------------------------------------------------------------------------------
/gateway/webssh.go:
--------------------------------------------------------------------------------
1 | /*
2 | * @Copyright Reserved By Janusec (https://www.janusec.com/).
3 | * @Author: U2
4 | * @Date: 2020-02-10 22:07:47
5 | * @Last Modified: U2, 2020-02-10 22:07:47
6 | */
7 |
8 | package gateway
9 |
10 | import (
11 | "bytes"
12 | "encoding/json"
13 | "errors"
14 | "io"
15 | "log"
16 | "net/http"
17 | "time"
18 |
19 | "janusec/data"
20 |
21 | "janusec/usermgmt"
22 | "janusec/utils"
23 |
24 | "github.com/gorilla/websocket"
25 | "golang.org/x/crypto/ssh"
26 | )
27 |
28 | // HostInfo : the information of remote Host
29 | type HostInfo struct {
30 | IP string `json:"ip"`
31 | Port string `json:"port"`
32 | Username string `json:"username"`
33 | Password string `json:"password"`
34 | }
35 |
36 | // SSH build connection
37 | func SSH(sshInput *io.WriteCloser, sshOutput *io.Reader, host *HostInfo, errChan chan<- error) {
38 | sshClient, err := ssh.Dial("tcp", host.IP+":"+host.Port, &ssh.ClientConfig{
39 | User: host.Username,
40 | Auth: []ssh.AuthMethod{ssh.Password(host.Password)},
41 | HostKeyCallback: ssh.InsecureIgnoreHostKey(),
42 | })
43 | if err != nil {
44 | errChan <- err
45 | utils.DebugPrintln("errChan", err)
46 | return
47 | }
48 | sshSession, err := sshClient.NewSession()
49 | if err != nil {
50 | utils.DebugPrintln("new ssh session", err)
51 | }
52 | defer sshSession.Close()
53 | *sshInput, err = sshSession.StdinPipe()
54 | if err != nil {
55 | utils.DebugPrintln("sshInput", err)
56 | }
57 | *sshOutput, err = sshSession.StdoutPipe()
58 | if err != nil {
59 | utils.DebugPrintln("sshOuput", err)
60 | }
61 | errChan <- err
62 | modes := ssh.TerminalModes{
63 | ssh.ECHO: 1,
64 | ssh.TTY_OP_ISPEED: 14400,
65 | ssh.TTY_OP_OSPEED: 14400,
66 | }
67 | err = sshSession.RequestPty("xterm", 40, 120, modes)
68 | if err != nil {
69 | utils.DebugPrintln("request pty", err)
70 | }
71 | err = sshSession.Shell()
72 | if err != nil {
73 | utils.DebugPrintln("start shell", err)
74 | }
75 | err = sshSession.Wait()
76 | errChan <- err
77 | }
78 |
79 | // RoutineOutput update the console display
80 | func RoutineOutput(outputTicker *time.Ticker, wsConn *websocket.Conn, sshOutput *io.Reader) {
81 | for range outputTicker.C {
82 | cmdOutput := make([]byte, 1024*10)
83 | n, err := (*sshOutput).Read(cmdOutput)
84 | if err != nil {
85 | // EOF
86 | return
87 | }
88 | if n > 0 {
89 | err := wsConn.WriteMessage(websocket.TextMessage, cmdOutput)
90 | if err != nil {
91 | return
92 | }
93 | }
94 | }
95 | }
96 |
97 | // WebSSHHandlerFunc Handle Web SSH
98 | func WebSSHHandlerFunc(w http.ResponseWriter, r *http.Request) {
99 | var isLogin bool
100 | isLogin, _ = usermgmt.IsLogIn(w, r)
101 | if !isLogin {
102 | GenResponseByObject(w, nil, errors.New("please login"))
103 | return
104 | }
105 | username := usermgmt.GetLoginUsername(r)
106 | var sshInput io.WriteCloser
107 | var sshOutput io.Reader //bytes.Buffer
108 | upgrader := websocket.Upgrader{
109 | ReadBufferSize: 1024,
110 | WriteBufferSize: 1024 * 10,
111 | CheckOrigin: func(r *http.Request) bool {
112 | return true
113 | },
114 | }
115 | wsConn, err := upgrader.Upgrade(w, r, nil)
116 | // websocket.Upgrade deprecated, add upgrader.Upgrade above, v1.2.0
117 | // wsConn, err := websocket.Upgrade(w, r, nil, 1024, 1024*10)
118 | if err != nil {
119 | log.Println("upgrade:", err)
120 | return
121 | }
122 | defer wsConn.Close()
123 | // Read SSH Parameters
124 | _, msg, err2 := wsConn.ReadMessage()
125 | if err2 != nil {
126 | utils.DebugPrintln("ReadMessage SSH Parameters Error:", err2)
127 | return
128 | }
129 | if !data.PrimarySetting.WebSSHEnabled {
130 | err = wsConn.WriteMessage(websocket.TextMessage, []byte("WebSSH disabled in settings!\r\n"))
131 | if err != nil {
132 | utils.DebugPrintln("WebSSHHandlerFunc wsConn.WriteMessage error", err)
133 | }
134 | return
135 | }
136 | var host HostInfo
137 | err = json.Unmarshal(msg, &host)
138 | if err != nil {
139 | utils.DebugPrintln("WebSSHHandlerFunc json.Unmarshal error", err)
140 | }
141 | if err = wsConn.WriteMessage(websocket.TextMessage, []byte("Connecting "+host.IP+":"+host.Port+" ... Please wait a moment!\r\n")); err != nil {
142 | return
143 | }
144 | errChan := make(chan error)
145 | go SSH(&sshInput, &sshOutput, &host, errChan)
146 | err = <-errChan
147 | if err != nil {
148 | err2 := wsConn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
149 | if err2 != nil {
150 | utils.DebugPrintln("WebSSHHandlerFunc wsConn.WriteMessage error", err2)
151 | }
152 | return
153 | }
154 | var logBuf bytes.Buffer
155 | outputTicker := time.NewTicker(100 * time.Millisecond)
156 | go RoutineOutput(outputTicker, wsConn, &sshOutput)
157 | for {
158 | select {
159 | case <-errChan:
160 | err2 := wsConn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
161 | if err2 != nil {
162 | utils.DebugPrintln("WebSSHHandlerFunc wsConn.WriteMessage error", err2)
163 | }
164 | return
165 | default:
166 | if wsConn == nil {
167 | return
168 | }
169 | _, msg, err2 := wsConn.ReadMessage()
170 | if err2 != nil {
171 | return
172 | }
173 | //log.Printf("Received: %s %v\n", string(msg), msg)
174 | if sshInput != nil {
175 | go CmdLog(&logBuf, username, &host, &msg)
176 | if _, err = sshInput.Write(msg); err != nil {
177 | return
178 | }
179 | }
180 | }
181 | }
182 | }
183 |
184 | // CmdLog write to log files
185 | func CmdLog(logBuf *bytes.Buffer, username string, host *HostInfo, cmdChars *[]byte) {
186 | for i := 0; i < len(*cmdChars); i++ {
187 | cmdChar := (*cmdChars)[i]
188 | switch cmdChar {
189 | case '\r', '\n':
190 | cmdStr := logBuf.String()
191 | hostInfo := host.Username + "@" + host.IP + ":" + host.Port
192 | utils.DebugPrintln("WebSSH User:", username, hostInfo, "Command:", cmdStr)
193 | logBuf.Reset()
194 | default:
195 | logBuf.WriteByte(cmdChar)
196 | }
197 | }
198 | }
199 |
--------------------------------------------------------------------------------