diff --git a/code/client/rule/vnc/process/process.go b/code/client/rule/vnc/process/process.go index 722b7fe..7560048 100644 --- a/code/client/rule/vnc/process/process.go +++ b/code/client/rule/vnc/process/process.go @@ -3,12 +3,15 @@ package process import ( "errors" "fmt" + "image" + "image/jpeg" "natpass/code/client/rule/vnc/vncnetwork" "natpass/code/utils" "net" "net/http" "os" "sync" + "time" "github.com/gorilla/websocket" "github.com/lwch/logging" @@ -112,3 +115,61 @@ func (p *Process) kill() { ps.Kill() } } + +// Close close process +func (p *Process) Close() { + if p.srv != nil { + p.srv.Close() + } + if p.chImage != nil { + close(p.chImage) + p.chImage = nil + } + if p.chClipboard != nil { + close(p.chClipboard) + p.chClipboard = nil + } + if p.chWrite != nil { + close(p.chWrite) + p.chWrite = nil + } + p.kill() +} + +// Capture capture desktop image +func (p *Process) Capture(timeout time.Duration) (*image.RGBA, error) { + var msg vncnetwork.VncMsg + msg.XType = vncnetwork.VncMsg_capture_req + p.chWrite <- &msg + trans := func(data *vncnetwork.ImageData) *image.RGBA { + img := image.NewRGBA(image.Rect(0, 0, int(data.GetWidth()), int(data.GetHeight()))) + copy(img.Pix, data.GetData()) + // dumpImage(img) + return img + } + if timeout > 0 { + select { + case data := <-p.chImage: + return trans(data), nil + case <-time.After(timeout): + return nil, errors.New("timeout") + } + } else { + data := <-p.chImage + return trans(data), nil + } +} + +func dumpImage(img image.Image) { + f, err := os.Create(`C:\Users\lwch\Pictures\debug.jpeg`) + if err != nil { + logging.Error("debug: %v", err) + return + } + defer f.Close() + err = jpeg.Encode(f, img, nil) + if err != nil { + logging.Error("encode: %v", err) + return + } +} diff --git a/code/client/rule/vnc/process/process_darwin.go b/code/client/rule/vnc/process/process_darwin.go new file mode 100644 index 0000000..9cb3f2c --- /dev/null +++ b/code/client/rule/vnc/process/process_darwin.go @@ -0,0 +1,11 @@ +package process + +import ( + "image" + "time" +) + +// CreateWorker create worker process +func CreateWorker(name, confDir string, showCursor bool) (*Process, error) { + return nil, ErrNotSupported +} diff --git a/code/client/rule/vnc/process/process_linux.go b/code/client/rule/vnc/process/process_linux.go new file mode 100644 index 0000000..e043a0b --- /dev/null +++ b/code/client/rule/vnc/process/process_linux.go @@ -0,0 +1,34 @@ +package process + +import ( + "fmt" + "natpass/code/client/rule/vnc/vncnetwork" + "os" + "os/exec" +) + +// CreateWorker create worker process +func CreateWorker(name, confDir string, showCursor bool) (*Process, error) { + dir, err := os.Executable() + if err != nil { + return nil, err + } + var p Process + p.chWrite = make(chan *vncnetwork.VncMsg) + p.chImage = make(chan *vncnetwork.ImageData) + p.chClipboard = make(chan *vncnetwork.ClipboardData) + port, err := p.listenAndServe() + if err != nil { + return nil, err + } + cmd := exec.Command(dir, "-conf", confDir, + "-action", "vnc.worker", + "-name", name, + "-vport", fmt.Sprintf("%d", port)) + err = cmd.Start() + if err != nil { + p.Close() + return nil, err + } + return &p, nil +} diff --git a/code/client/rule/vnc/process/process_windows.go b/code/client/rule/vnc/process/process_windows.go index 9f6204d..0cf02cb 100644 --- a/code/client/rule/vnc/process/process_windows.go +++ b/code/client/rule/vnc/process/process_windows.go @@ -123,61 +123,3 @@ func createWorker(name, confDir string, tk windows.Token, showCursor bool) (*Pro p.pid = int(process.ProcessId) return &p, nil } - -// Close close process -func (p *Process) Close() { - if p.srv != nil { - p.srv.Close() - } - if p.chImage != nil { - close(p.chImage) - p.chImage = nil - } - if p.chClipboard != nil { - close(p.chClipboard) - p.chClipboard = nil - } - if p.chWrite != nil { - close(p.chWrite) - p.chWrite = nil - } - p.kill() -} - -// Capture capture desktop image -func (p *Process) Capture(timeout time.Duration) (*image.RGBA, error) { - var msg vncnetwork.VncMsg - msg.XType = vncnetwork.VncMsg_capture_req - p.chWrite <- &msg - trans := func(data *vncnetwork.ImageData) *image.RGBA { - img := image.NewRGBA(image.Rect(0, 0, int(data.GetWidth()), int(data.GetHeight()))) - copy(img.Pix, data.GetData()) - // dumpImage(img) - return img - } - if timeout > 0 { - select { - case data := <-p.chImage: - return trans(data), nil - case <-time.After(timeout): - return nil, errors.New("timeout") - } - } else { - data := <-p.chImage - return trans(data), nil - } -} - -func dumpImage(img image.Image) { - f, err := os.Create(`C:\Users\lwch\Pictures\debug.jpeg`) - if err != nil { - logging.Error("debug: %v", err) - return - } - defer f.Close() - err = jpeg.Encode(f, img, nil) - if err != nil { - logging.Error("encode: %v", err) - return - } -} diff --git a/code/client/rule/vnc/process/process_xx.go b/code/client/rule/vnc/process/process_xx.go deleted file mode 100644 index b2e6eaf..0000000 --- a/code/client/rule/vnc/process/process_xx.go +++ /dev/null @@ -1,23 +0,0 @@ -//go:build !windows -// +build !windows - -package process - -import ( - "image" - "time" -) - -// CreateWorker create worker process -func CreateWorker(name, confDir string, showCursor bool) (*Process, error) { - return nil, ErrNotSupported -} - -// Close close process -func (p *Process) Close() { -} - -// Capture capture desktop image -func (p *Process) Capture(timeout time.Duration) (*image.RGBA, error) { - return nil, nil -}