package main

import (
	"fmt"
	"io"
	"os"
	"os/exec"

	"github.com/creack/pty"
	"github.com/gliderlabs/ssh"
)

func shell(s ssh.Session) {
	ptyReq, winCh, isPty := s.Pty()

	if !isPty {
		fmt.Fprintln(s, "Must be PTY")
		s.Exit(1)
		return
	}

	shell := os.Getenv("SHELL")
	if shell == "" {
		shell = "/bin/sh"
	}

	cmd := exec.Command(shell)
	cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))

	ptmx, err := pty.Start(cmd)
	if err != nil {
		fmt.Fprintln(s, "Can not start shell")
		fmt.Println(err)
		s.Exit(1)
		return
	}
	defer func() { _ = ptmx.Close() }()

	go func() {
		for win := range winCh {
			pty.Setsize(ptmx, &pty.Winsize{
				Rows: uint16(win.Height),
				Cols: uint16(win.Width),
			})
		}
	}()

	go func() {
		io.Copy(s, ptmx)
		s.Close()
	}()

	go func() {
		io.Copy(ptmx, s)
		ptmx.Close()
	}()

	if err := cmd.Wait(); err != nil {
		s.Exit(1)
	} else {
		s.Exit(0)
	}
}