kanshan/megrez
1
0
Fork 0
mirror of https://github.com/XShengTech/MEGREZ.git synced 2026-01-14 09:07:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

[Perf] 🚀 Set Jupter Token for Safety Improve #7

Browse Source
This commit is contained in:
Harry-zklcdc 2025-01-13 19:06:51 +08:00
parent 35f2b99b17
commit 9daa535c9c
4 changed files with 93 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
services/instanceController/create.go
View File

@ -47,6 +47,8 @@ func Create(instance *models.Instances) (containerName, volumeName string, err e
return "", "", err
}
go SetJupterPassword(server.IP, server.Port, server.Apikey, containerName, instance.SshPasswd)
portBindings, err := GetPortForward(server.IP, server.Port, server.Apikey, containerName)
if err != nil {
deleteInstance(server.IP, server.Port, server.Apikey, containerName)

85
services/instanceController/execute.go
View File

@ -56,3 +56,88 @@ func SetRootPassword(ip string, port int, apikey string,
return nil
}
func SetJupterPassword(ip string, port int, apikey string,
containerName, password string) (err error) {
l.SetFunction("SetJupterPassword")
// Set Jupyter Password
data := executeReq{
Cmd: []string{
"sed",
"-i",
"/^c.ServerApp.token = ''/c\\c.ServerApp.token = '" + password + "'",
"/root/.jupyter/jupyter_notebook_config.py",
},
}
reqBytes, err := json.Marshal(data)
if err != nil {
l.Error("marshal request data error: %v", err)
return
}
c := request.NewRequest().Post().
SetUrl("http://" + ip + ":" + strconv.Itoa(port) + apiPrefix + instancePrefix + "/" + containerName + instanceExecute).
SetAuthorization("Bearer " + apikey).
SetUserAgent("megrez").
SetBody(bytes.NewBuffer(reqBytes))
c.Do()
if c.GetStatusCode() != 200 {
l.Error("set jupter password error: %d", c.GetStatusCode())
return errors.New("set jupter password request error")
}
var res resStruct
err = json.Unmarshal(c.GetBody(), &res)
if err != nil {
l.Error("unmarshal response data error: %v", err)
return err
}
if res.Code != 200 {
l.Error("set jupter password code: %d, error: %s", res.Code, res.Msg)
return errors.New(res.Msg)
}
// Restart Jupyter
data = executeReq{
Cmd: []string{
"service",
"jupyter",
"restart",
},
}
reqBytes, err = json.Marshal(data)
if err != nil {
l.Error("marshal request data error: %v", err)
return
}
c = request.NewRequest().Post().
SetUrl("http://" + ip + ":" + strconv.Itoa(port) + apiPrefix + instancePrefix + "/" + containerName + instanceExecute).
SetAuthorization("Bearer " + apikey).
SetUserAgent("megrez").
SetBody(bytes.NewBuffer(reqBytes))
c.Do()
if c.GetStatusCode() != 200 {
l.Error("restart jupter error: %d", c.GetStatusCode())
return errors.New("restart jupter request error")
}
err = json.Unmarshal(c.GetBody(), &res)
if err != nil {
l.Error("unmarshal response data error: %v", err)
return err
}
if res.Code != 200 {
l.Error("restart jupter code: %d, error: %s", res.Code, res.Msg)
return errors.New(res.Msg)
}
return nil
}

4
services/instanceController/patch.go
View File

@ -95,6 +95,8 @@ func Patch(instance *models.Instances, gpuCount, volumeSize int, cpuOnly bool) (
return err
}
go SetJupterPassword(server.IP, server.Port, server.Apikey, instance.ContainerName, instance.SshPasswd)
portBindings, err := GetPortForward(server.IP, server.Port, server.Apikey, instance.ContainerName)
if err != nil {
l.Error("get port forward error: %v", err)
@ -138,6 +140,8 @@ func Patch(instance *models.Instances, gpuCount, volumeSize int, cpuOnly bool) (
return err
}
go SetJupterPassword(server.IP, server.Port, server.Apikey, instance.ContainerName, instance.SshPasswd)
portBindings, err := GetPortForward(server.IP, server.Port, server.Apikey, instance.ContainerName)
if err != nil {
l.Error("get port forward error: %v", err)

2
services/instanceController/restart.go
View File

@ -48,6 +48,8 @@ func Restart(instance *models.Instances) (err error) {
return err
}
go SetJupterPassword(server.IP, server.Port, server.Apikey, instance.ContainerName, instance.SshPasswd)
portBindings, err := GetPortForward(server.IP, server.Port, server.Apikey, instance.ContainerName)
if err != nil {
l.Error("get port forward error: %v", err)