diff --git a/main.go b/main.go
index 932d347..9874217 100644
--- a/main.go
+++ b/main.go
@@ -146,6 +146,8 @@ Flags:
 	bufio.NewReader(os.Stdin).ReadString('\n')
 
 	ptyMaster := ptyMasterNew()
+	defer ptyMaster.Restore()
+
 	ptyMaster.Start(*commandName, strings.Fields(*commandArgs))
 
 	var pty server.PTYHandler = ptyMaster
diff --git a/pty_master.go b/pty_master.go
index b39366b..8a26a09 100644
--- a/pty_master.go
+++ b/pty_master.go
@@ -101,7 +101,10 @@ func (pty *ptyMaster) Refresh() {
 
 func (pty *ptyMaster) Wait() (err error) {
 	err = pty.command.Wait()
-	// The terminal has to be restored from the RAW state, to its initial state
+	return
+}
+
+func (pty *ptyMaster) Restore() {
 	terminal.Restore(0, pty.terminalInitState)
 	return
 }