From 12b0db07daa7098b0df9d1de1bd53bfc288fb4b8 Mon Sep 17 00:00:00 2001
From: Manuel <manuel.sabban@nbs-system.com>
Date: Tue, 2 Jan 2018 16:31:34 +0100
Subject: [PATCH] add audit feature.

---
 main.go                              |  7 ++++++-
 pkg/bastionsession/bastionsession.go | 24 ++++++++++++++++++------
 ssh.go                               |  2 ++
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/main.go b/main.go
index 5cc3bb3..1a709e8 100644
--- a/main.go
+++ b/main.go
@@ -68,6 +68,11 @@ func main() {
 					Name:  "aes-key",
 					Usage: "Encrypt sensitive data in database (length: 16, 24 or 32)",
 				},
+				cli.StringFlag{
+					Name: "logs-location",
+					Value: "/var/log/sshportal",
+					Usage: "Store user session files",
+				},
 			},
 		}, {
 			Name:   "healthcheck",
@@ -125,7 +130,7 @@ func server(c *cli.Context) error {
 	opts = append(opts, ssh.PublicKeyAuth(publicKeyAuthHandler(db, c)))
 	opts = append(opts, ssh.PasswordAuth(passwordAuthHandler(db, c)))
 
-	// retrieve sshportal SSH private key from databse
+	// retrieve sshportal SSH private key from database
 	opts = append(opts, func(srv *ssh.Server) error {
 		var key SSHKey
 		if err = SSHKeysByIdentifiers(db, []string{"host"}).First(&key).Error; err != nil {
diff --git a/pkg/bastionsession/bastionsession.go b/pkg/bastionsession/bastionsession.go
index f9bc417..4acef03 100644
--- a/pkg/bastionsession/bastionsession.go
+++ b/pkg/bastionsession/bastionsession.go
@@ -3,13 +3,18 @@ package bastionsession
 import (
 	"errors"
 	"io"
-
+	"strings"
+	"time"
+	"os"
+	
 	"github.com/gliderlabs/ssh"
+	"github.com/sabban/sshportal/pkg/logchannel"
 	gossh "golang.org/x/crypto/ssh"
 )
 
 type Config struct {
 	Addr         string
+	Logs         string
 	ClientConfig *gossh.ClientConfig
 }
 
@@ -35,21 +40,28 @@ func ChannelHandler(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewCh
 	if err != nil {
 		return err
 	}
-
+	user := conn.User()
 	// pipe everything
-	return pipe(lreqs, rreqs, lch, rch)
+	return pipe(lreqs, rreqs, lch, rch, config.Logs, user)
 }
 
-func pipe(lreqs, rreqs <-chan *gossh.Request, lch, rch gossh.Channel) error {
+func pipe(lreqs, rreqs <-chan *gossh.Request, lch, rch gossh.Channel, logs_location string, user string) error {
 	defer func() {
 		_ = lch.Close()
 		_ = rch.Close()
 	}()
 
 	errch := make(chan error, 1)
-
+	file_name := strings.Join([]string{logs_location, "/", user, "-", time.Now().Format("RFC3339")}, "") // get user
+	f, err := os.OpenFile(file_name, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0640)
+	if err != nil {
+		errch <- errors.New("Opening session file" + file_name + "failed.")
+	}
+	defer f.Close()
+	wrappedlch := logchannel.New(lch, f) 
+	
 	go func() {
-		_, _ = io.Copy(lch, rch)
+		_, _ = io.Copy(wrappedlch, rch)
 		errch <- errors.New("lch closed the connection")
 	}()
 
diff --git a/ssh.go b/ssh.go
index 5076964..099a329 100644
--- a/ssh.go
+++ b/ssh.go
@@ -96,6 +96,7 @@ func channelHandler(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewCh
 	}
 
 	actx := ctx.Value(authContextKey).(*authContext)
+	logs_locations := ctx.Value("logs-location").(*authContext)
 
 	switch actx.userType() {
 	case UserTypeBastion:
@@ -129,6 +130,7 @@ func channelHandler(srv *ssh.Server, conn *gossh.ServerConn, newChan gossh.NewCh
 
 		err = bastionsession.ChannelHandler(srv, conn, newChan, ctx, bastionsession.Config{
 			Addr:         host.Addr,
+			Logs:         logs_location,
 			ClientConfig: clientConfig,
 		})