package main

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/binary"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"log"
	"net"

	"git.snorba.art/hugo/nssh/commands"
	"git.snorba.art/hugo/nssh/storage"
	"golang.org/x/crypto/ssh"
)

func generatePrivateKey(bitSize int) (*rsa.PrivateKey, error) {
	// Private Key generation
	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
	if err != nil {
		return nil, err
	}

	// Validate Private Key
	err = privateKey.Validate()
	if err != nil {
		return nil, err
	}

	log.Println("Private Key generated")
	return privateKey, nil
}

func encodePrivateKeyToPEM(privateKey *rsa.PrivateKey) []byte {
	// Get ASN.1 DER format
	privDER := x509.MarshalPKCS1PrivateKey(privateKey)

	// pem.Block
	privBlock := pem.Block{
		Type:    "RSA PRIVATE KEY",
		Headers: nil,
		Bytes:   privDER,
	}

	// Private key in PEM format
	privatePEM := pem.EncodeToMemory(&privBlock)

	return privatePEM
}

// parseDims extracts two uint32s from the provided buffer.
func parseDims(b []byte) (uint32, uint32) {
	w := binary.BigEndian.Uint32(b)
	h := binary.BigEndian.Uint32(b[4:])
	return w, h
}

type Sshd struct {
	AuthorizedKeysMap map[string]bool
	Filestore         storage.Filestore
}

func (s *Sshd) Listen(iface string) error {
	// An SSH server is represented by a ServerConfig, which holds
	// certificate details and handles authentication of ServerConns.
	config := &ssh.ServerConfig{
		// Remove to disable public key auth.
		PublicKeyCallback: func(c ssh.ConnMetadata, pubKey ssh.PublicKey) (*ssh.Permissions, error) {
			if s.AuthorizedKeysMap[string(pubKey.Marshal())] {
				return &ssh.Permissions{
					// Record the public key used for authentication.
					Extensions: map[string]string{
						"pubkey-fp": ssh.FingerprintSHA256(pubKey),
					},
				}, nil
			}
			return nil, fmt.Errorf("unknown public key for %q", c.User())
		},
	}

	bitSize := 4096

	privateKey, err := generatePrivateKey(bitSize)
	if err != nil {
		return err
	}

	privateBytes := encodePrivateKeyToPEM(privateKey)

	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		return fmt.Errorf("Failed to parse private key: %w", err)
	}

	config.AddHostKey(private)

	// Once a ServerConfig has been configured, connections can be
	// accepted.
	listener, err := net.Listen("tcp", iface)
	if err != nil {
		return err
	}

	return s.HandleRequests(listener, config)
}

func (s *Sshd) HandleRequests(listener net.Listener, config *ssh.ServerConfig) error {
	for {
		nConn, err := listener.Accept()
		if err != nil {
			return fmt.Errorf("failed to accept incoming connection: %w", err)
		}

		// Before use, a handshake must be performed on the incoming
		// net.Conn.
		conn, chans, reqs, err := ssh.NewServerConn(nConn, config)
		if err != nil {
			log.Printf("failed to handshake: %s", err)
			continue
		}
		log.Printf("logged in with key %s", conn.Permissions.Extensions["pubkey-fp"])

		// The incoming Request channel must be serviced.
		go ssh.DiscardRequests(reqs)

		// Service the incoming Channel channel.
		for newChannel := range chans {
			// Channels have a type, depending on the application level
			// protocol intended. In the case of a shell, the type is
			// "session" and ServerShell may be used to present a simple
			// terminal interface.
			if newChannel.ChannelType() != "session" {
				newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
				continue
			}
			channel, requests, err := newChannel.Accept()
			if err != nil {
				return fmt.Errorf("Could not accept channel: %w", err)
			}

			session := &SshSession{
				shell:     commands.NewShell("user@host:/# "),
				filestore: s.Filestore,
			}

			// Sessions have out-of-band requests such as "shell",
			// "pty-req" and "env".  Here we handle only the
			// "shell" request.
			go session.HandleRequests(requests)

			go session.RunShell(channel)
		}
	}
}

type SshSession struct {
	shell     *commands.Shell
	filestore storage.Filestore
}

func (s *SshSession) HandleRequests(in <-chan *ssh.Request) {
	for req := range in {
		log.Printf("received request of type \"%s\"", req.Type)

		var ok bool
		switch req.Type {
		case "shell":
			ok = true
		case "pty-req":
			// Responding 'ok' here will let the client
			// know we have a pty ready for input
			ok = true
			// Parse body...
			termLen := req.Payload[3]
			termEnv := string(req.Payload[4 : termLen+4])
			w, h := parseDims(req.Payload[termLen+4:])
			s.shell.SetWinsize(w, h)
			log.Printf("pty-req '%s'", termEnv)
		}

		req.Reply(ok, nil)
	}
}

func (s *SshSession) RunShell(channel io.ReadWriteCloser) {
	s.shell.Attach(channel)
	handler := commands.NewCommandHandler(s.shell, s.filestore)

	go func() {
		for {
			comm, err := s.shell.ReadCommand()
			if err != nil {
				if errors.Is(err, io.EOF) {
					log.Println("Received EOF, closing TTY")
				} else {
					log.Println("Error reading command: ", err)
				}

				if !errors.Is(err, commands.ErrEmptyCommand) {
					s.shell.Close()
					return
				}
			} else {
				log.Printf("command: \"%s\" %s", comm.Name, comm.Arguments)

				err := handler.Handle(comm)
				if err != nil {
					log.Println("Error handling command: " + err.Error())
				}
			}

			err = s.shell.WritePrompt()
			if err != nil {
				log.Println(err)
			}
		}
	}()
}