package main

import (
	"archive/zip"
	"bytes"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"math"
	"math/big"
	"net"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
)

type tlsGenerateRequest struct {
	CommonName         string `json:"commonName"`
	Organization       string `json:"organization"`
	OrganizationalUnit string `json:"organizationalUnit"`
	Locality           string `json:"locality"`
	State              string `json:"state"`
	Country            string `json:"country"`
	DNSNames           string `json:"dnsNames"`
	ValidDays          int    `json:"validDays"`
	KeySize            int    `json:"keySize"`
}

type dnsLookupRequest struct {
	Host string `json:"host"`
}

type tslCheckRequest struct {
	URL string `json:"url"`
}

type pemCheckRequest struct {
	PEM string `json:"pem"`
}

type apiError struct {
	Error string `json:"error"`
}

func main() {
	port := getenv("PORT", "8080")
	certFile := strings.TrimSpace(os.Getenv("TLS_CERT_FILE"))
	keyFile := strings.TrimSpace(os.Getenv("TLS_KEY_FILE"))

	router := newRouter()
	addr := ":" + port

	var err error
	if certFile != "" || keyFile != "" {
		if certFile == "" || keyFile == "" {
			panic("both TLS_CERT_FILE and TLS_KEY_FILE must be set to enable TLS")
		}
		err = router.RunTLS(addr, certFile, keyFile)
	} else {
		err = router.Run(addr)
	}

	if err != nil {
		panic(err)
	}
}

func newRouter() *gin.Engine {
	router := gin.Default()
	router.LoadHTMLGlob("templates/*")
	router.Static("/static", "./static")
	_ = router.SetTrustedProxies(nil)

	router.GET("/", func(c *gin.Context) {
		c.HTML(http.StatusOK, "index.html", gin.H{
			"title": "Network Tools",
		})
	})

	api := router.Group("/api")
	{
		api.POST("/tls/generate", handleTLSGenerate)
		api.POST("/dns/lookup", handleDNSLookup)
		api.POST("/tsl/check", handleTSLCheck)
		api.POST("/pem/check", handlePEMCheck)
	}

	return router
}

func handleTLSGenerate(c *gin.Context) {
	var req tlsGenerateRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, apiError{Error: "invalid request body"})
		return
	}

	commonName := strings.TrimSpace(req.CommonName)
	if commonName == "" {
		c.JSON(http.StatusBadRequest, apiError{Error: "commonName is required"})
		return
	}

	validDays := req.ValidDays
	if validDays <= 0 {
		validDays = 365
	}

	keySize := normalizeKeySize(req.KeySize)
	if keySize == 0 {
		c.JSON(http.StatusBadRequest, apiError{Error: "keySize must be one of 2048, 3072, or 4096"})
		return
	}

	dnsNames := splitCSV(req.DNSNames)
	if len(dnsNames) == 0 {
		dnsNames = []string{commonName}
	}

	subject := pkix.Name{
		CommonName:         commonName,
		Organization:       splitCSV(req.Organization),
		OrganizationalUnit: splitCSV(req.OrganizationalUnit),
		Locality:           splitCSV(req.Locality),
		Province:           splitCSV(req.State),
		Country:            splitCSV(req.Country),
	}

	certPEM, keyPEM, publicKeyPEM, csrPEM, parsed, err := generateCertificate(subject, dnsNames, validDays, keySize)
	if err != nil {
		c.JSON(http.StatusInternalServerError, apiError{Error: err.Error()})
		return
	}

	baseName := sanitizeFilename(commonName)
	zipBytes, err := buildTLSArchive(baseName, certPEM, keyPEM, publicKeyPEM, csrPEM)
	if err != nil {
		c.JSON(http.StatusInternalServerError, apiError{Error: err.Error()})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"certificatePem":           string(certPEM),
		"selfSignedCertificatePem": string(certPEM),
		"privateKeyPem":            string(keyPEM),
		"publicKeyPem":             string(publicKeyPEM),
		"csrPem":                   string(csrPEM),
		"subject":                  parsed.Subject.String(),
		"issuer":                   parsed.Issuer.String(),
		"serialNumber":             parsed.SerialNumber.String(),
		"notBefore":                parsed.NotBefore.Format(time.RFC3339),
		"notAfter":                 parsed.NotAfter.Format(time.RFC3339),
		"dnsNames":                 parsed.DNSNames,
		"keySize":                  keySize,
		"isSelfSigned":             parsed.Subject.String() == parsed.Issuer.String(),
		"organization":             parsed.Subject.Organization,
		"ou":                       parsed.Subject.OrganizationalUnit,
		"locality":                 parsed.Subject.Locality,
		"state":                    parsed.Subject.Province,
		"country":                  parsed.Subject.Country,
		"zipFilename":              baseName + ".zip",
		"zipBase64":                base64.StdEncoding.EncodeToString(zipBytes),
	})
}

func handleDNSLookup(c *gin.Context) {
	var req dnsLookupRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, apiError{Error: "invalid request body"})
		return
	}

	host := normalizeHost(req.Host)
	if host == "" {
		c.JSON(http.StatusBadRequest, apiError{Error: "host is required"})
		return
	}

	ips, ipErr := net.LookupIP(host)
	txt, txtErr := net.LookupTXT(host)
	mx, mxErr := net.LookupMX(host)
	ns, nsErr := net.LookupNS(host)
	cname, cnameErr := net.LookupCNAME(host)

	if ipErr != nil && txtErr != nil && mxErr != nil && nsErr != nil && cnameErr != nil {
		c.JSON(http.StatusBadGateway, apiError{Error: fmt.Sprintf("lookup failed for %s", host)})
		return
	}

	ipStrings := make([]string, 0, len(ips))
	for _, ip := range ips {
		ipStrings = append(ipStrings, ip.String())
	}

	txtStrings := make([]string, 0, len(txt))
	txtStrings = append(txtStrings, txt...)

	mxStrings := make([]gin.H, 0, len(mx))
	for _, record := range mx {
		mxStrings = append(mxStrings, gin.H{
			"host": record.Host,
			"pref": record.Pref,
		})
	}

	nsStrings := make([]string, 0, len(ns))
	for _, record := range ns {
		nsStrings = append(nsStrings, record.Host)
	}

	errorsByType := gin.H{}
	if ipErr != nil {
		errorsByType["a_aaaa"] = ipErr.Error()
	}
	if txtErr != nil {
		errorsByType["txt"] = txtErr.Error()
	}
	if mxErr != nil {
		errorsByType["mx"] = mxErr.Error()
	}
	if nsErr != nil {
		errorsByType["ns"] = nsErr.Error()
	}
	if cnameErr != nil {
		errorsByType["cname"] = cnameErr.Error()
	}

	c.JSON(http.StatusOK, gin.H{
		"host":   host,
		"ips":    ipStrings,
		"txt":    txtStrings,
		"mx":     mxStrings,
		"ns":     nsStrings,
		"cname":  strings.TrimSuffix(cname, "."),
		"errors": errorsByType,
	})
}

func handleTSLCheck(c *gin.Context) {
	var req tslCheckRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, apiError{Error: "invalid request body"})
		return
	}

	targetURL := strings.TrimSpace(req.URL)
	if targetURL == "" {
		c.JSON(http.StatusBadRequest, apiError{Error: "url is required"})
		return
	}

	parsedURL, err := normalizeURL(targetURL)
	if err != nil {
		c.JSON(http.StatusBadRequest, apiError{Error: err.Error()})
		return
	}

	host := parsedURL.Hostname()
	port := parsedURL.Port()
	if port == "" {
		port = "443"
	}

	conn, err := tls.DialWithDialer(
		&net.Dialer{Timeout: 10 * time.Second},
		"tcp",
		net.JoinHostPort(host, port),
		&tls.Config{ServerName: host},
	)
	if err != nil {
		c.JSON(http.StatusBadGateway, apiError{Error: err.Error()})
		return
	}
	defer conn.Close()

	state := conn.ConnectionState()
	if len(state.PeerCertificates) == 0 {
		c.JSON(http.StatusBadGateway, apiError{Error: "no peer certificates returned"})
		return
	}

	leaf := state.PeerCertificates[0]
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: leaf.Raw})

	chain := make([]gin.H, 0, len(state.PeerCertificates))
	for _, cert := range state.PeerCertificates {
		chain = append(chain, gin.H{
			"subject":      cert.Subject.String(),
			"issuer":       cert.Issuer.String(),
			"serialNumber": cert.SerialNumber.String(),
			"notBefore":    cert.NotBefore.Format(time.RFC3339),
			"notAfter":     cert.NotAfter.Format(time.RFC3339),
			"dnsNames":     cert.DNSNames,
		})
	}

	validDays := int(math.Ceil(leaf.NotAfter.Sub(leaf.NotBefore).Hours() / 24))
	if validDays < 1 {
		validDays = 1
	}

	keySize := publicKeySize(leaf.PublicKey)
	if keySize == 0 {
		keySize = 2048
	}

	c.JSON(http.StatusOK, gin.H{
		"url":                parsedURL.String(),
		"serverName":         host,
		"version":            tlsVersionName(state.Version),
		"cipherSuite":        tls.CipherSuiteName(state.CipherSuite),
		"negotiatedProtocol": state.NegotiatedProtocol,
		"certificatePem":     string(certPEM),
		"chain":              chain,
		"leafTemplate": gin.H{
			"commonName":         leaf.Subject.CommonName,
			"organization":       strings.Join(leaf.Subject.Organization, ", "),
			"organizationalUnit": strings.Join(leaf.Subject.OrganizationalUnit, ", "),
			"locality":           strings.Join(leaf.Subject.Locality, ", "),
			"state":              strings.Join(leaf.Subject.Province, ", "),
			"country":            strings.Join(leaf.Subject.Country, ", "),
			"dnsNames":           strings.Join(leaf.DNSNames, ", "),
			"validDays":          validDays,
			"keySize":            keySize,
		},
	})
}

func handlePEMCheck(c *gin.Context) {
	var req pemCheckRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, apiError{Error: "invalid request body"})
		return
	}

	input := strings.TrimSpace(req.PEM)
	if input == "" {
		c.JSON(http.StatusBadRequest, apiError{Error: "pem is required"})
		return
	}

	blocks, err := inspectPEM(input)
	if err != nil {
		c.JSON(http.StatusBadRequest, apiError{Error: err.Error()})
		return
	}

	c.JSON(http.StatusOK, gin.H{
		"blocks": blocks,
		"count":  len(blocks),
	})
}

func generateCertificate(subject pkix.Name, dnsNames []string, validDays int, keySize int) ([]byte, []byte, []byte, []byte, *x509.Certificate, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, keySize)
	if err != nil {
		return nil, nil, nil, nil, nil, err
	}

	serialLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialLimit)
	if err != nil {
		return nil, nil, nil, nil, nil, err
	}

	notBefore := time.Now().UTC()
	template := &x509.Certificate{
		SerialNumber:          serialNumber,
		Subject:               subject,
		NotBefore:             notBefore,
		NotAfter:              notBefore.Add(time.Duration(validDays) * 24 * time.Hour),
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth, x509.ExtKeyUsageClientAuth},
		BasicConstraintsValid: true,
		DNSNames:              dnsNames,
	}

	derBytes, err := x509.CreateCertificate(rand.Reader, template, template, &privateKey.PublicKey, privateKey)
	if err != nil {
		return nil, nil, nil, nil, nil, err
	}

	csrBytes, err := x509.CreateCertificateRequest(rand.Reader, &x509.CertificateRequest{
		Subject:            subject,
		SignatureAlgorithm: x509.SHA256WithRSA,
		DNSNames:           dnsNames,
	}, privateKey)
	if err != nil {
		return nil, nil, nil, nil, nil, err
	}

	publicKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, nil, nil, nil, nil, err
	}

	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)})
	publicKeyPEM := pem.EncodeToMemory(&pem.Block{Type: "PUBLIC KEY", Bytes: publicKeyBytes})
	csrPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: csrBytes})

	parsedCert, err := x509.ParseCertificate(derBytes)
	if err != nil {
		return nil, nil, nil, nil, nil, err
	}

	return certPEM, keyPEM, publicKeyPEM, csrPEM, parsedCert, nil
}

func buildTLSArchive(baseName string, certPEM []byte, keyPEM []byte, publicKeyPEM []byte, csrPEM []byte) ([]byte, error) {
	var buffer bytes.Buffer
	archive := zip.NewWriter(&buffer)

	files := []struct {
		name    string
		content []byte
	}{
		{name: baseName + ".crt.pem", content: certPEM},
		{name: baseName + ".key.pem", content: keyPEM},
		{name: baseName + ".pub.pem", content: publicKeyPEM},
		{name: baseName + ".csr.pem", content: csrPEM},
	}

	for _, file := range files {
		writer, err := archive.Create(file.name)
		if err != nil {
			return nil, err
		}
		if _, err := writer.Write(file.content); err != nil {
			return nil, err
		}
	}

	if err := archive.Close(); err != nil {
		return nil, err
	}

	return buffer.Bytes(), nil
}

func inspectPEM(input string) ([]gin.H, error) {
	var blocks []gin.H
	remaining := []byte(input)

	for len(strings.TrimSpace(string(remaining))) > 0 {
		block, rest := pem.Decode(remaining)
		if block == nil {
			return nil, errors.New("unable to decode PEM data")
		}

		inspected := gin.H{
			"pemType": block.Type,
		}

		switch block.Type {
		case "CERTIFICATE":
			cert, err := x509.ParseCertificate(block.Bytes)
			if err != nil {
				return nil, err
			}
			inspected["kind"] = "certificate"
			inspected["subject"] = cert.Subject.String()
			inspected["issuer"] = cert.Issuer.String()
			inspected["serialNumber"] = cert.SerialNumber.String()
			inspected["notBefore"] = cert.NotBefore.Format(time.RFC3339)
			inspected["notAfter"] = cert.NotAfter.Format(time.RFC3339)
			inspected["dnsNames"] = cert.DNSNames
			inspected["isCA"] = cert.IsCA
			inspected["signatureAlgorithm"] = cert.SignatureAlgorithm.String()
			inspected["publicKeyAlgorithm"] = cert.PublicKeyAlgorithm.String()
		case "CERTIFICATE REQUEST", "NEW CERTIFICATE REQUEST":
			csr, err := x509.ParseCertificateRequest(block.Bytes)
			if err != nil {
				return nil, err
			}
			inspected["kind"] = "certificate_request"
			inspected["subject"] = csr.Subject.String()
			inspected["dnsNames"] = csr.DNSNames
			inspected["emailAddresses"] = csr.EmailAddresses
			inspected["signatureAlgorithm"] = csr.SignatureAlgorithm.String()
			inspected["publicKeyAlgorithm"] = csr.PublicKeyAlgorithm.String()
		default:
			keyInfo, err := inspectKeyBlock(block)
			if err != nil {
				return nil, err
			}
			for key, value := range keyInfo {
				inspected[key] = value
			}
		}

		blocks = append(blocks, inspected)
		remaining = rest
	}

	if len(blocks) == 0 {
		return nil, errors.New("no PEM blocks found")
	}

	return blocks, nil
}

func inspectKeyBlock(block *pem.Block) (gin.H, error) {
	switch block.Type {
	case "RSA PRIVATE KEY":
		key, err := x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		return gin.H{
			"kind":      "private_key",
			"algorithm": "RSA",
			"size":      key.N.BitLen(),
		}, nil
	case "EC PRIVATE KEY":
		key, err := x509.ParseECPrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		return gin.H{
			"kind":      "private_key",
			"algorithm": "EC",
			"curve":     key.Curve.Params().Name,
		}, nil
	case "PRIVATE KEY":
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		return describePrivateKey(key), nil
	case "PUBLIC KEY":
		key, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, err
		}
		return describePublicKey(key), nil
	default:
		return nil, fmt.Errorf("unsupported PEM block type: %s", block.Type)
	}
}

func describePrivateKey(key any) gin.H {
	switch typed := key.(type) {
	case *rsa.PrivateKey:
		return gin.H{
			"kind":      "private_key",
			"algorithm": "RSA",
			"size":      typed.N.BitLen(),
		}
	default:
		return gin.H{
			"kind":      "private_key",
			"algorithm": reflect.TypeOf(key).String(),
		}
	}
}

func describePublicKey(key any) gin.H {
	switch typed := key.(type) {
	case *rsa.PublicKey:
		return gin.H{
			"kind":      "public_key",
			"algorithm": "RSA",
			"size":      typed.N.BitLen(),
		}
	default:
		return gin.H{
			"kind":      "public_key",
			"algorithm": reflect.TypeOf(key).String(),
		}
	}
}

func splitCSV(input string) []string {
	parts := strings.Split(input, ",")
	values := make([]string, 0, len(parts))
	for _, part := range parts {
		trimmed := strings.TrimSpace(part)
		if trimmed != "" {
			values = append(values, trimmed)
		}
	}
	return values
}

func normalizeHost(input string) string {
	value := strings.TrimSpace(input)
	value = strings.TrimPrefix(value, "https://")
	value = strings.TrimPrefix(value, "http://")
	value = strings.Trim(value, "/")
	if host, _, err := net.SplitHostPort(value); err == nil {
		return host
	}
	if parsed, err := url.Parse("https://" + value); err == nil {
		return parsed.Hostname()
	}
	return value
}

func normalizeURL(input string) (*url.URL, error) {
	value := strings.TrimSpace(input)
	if value == "" {
		return nil, errors.New("url is required")
	}
	if !strings.Contains(value, "://") {
		value = "https://" + value
	}
	parsed, err := url.Parse(value)
	if err != nil {
		return nil, err
	}
	if parsed.Hostname() == "" {
		return nil, errors.New("url must include a hostname")
	}
	return parsed, nil
}

func tlsVersionName(version uint16) string {
	switch version {
	case tls.VersionTLS10:
		return "TLS 1.0"
	case tls.VersionTLS11:
		return "TLS 1.1"
	case tls.VersionTLS12:
		return "TLS 1.2"
	case tls.VersionTLS13:
		return "TLS 1.3"
	default:
		return fmt.Sprintf("0x%x", version)
	}
}

func normalizeKeySize(value int) int {
	switch value {
	case 0, 2048:
		return 2048
	case 3072:
		return 3072
	case 4096:
		return 4096
	default:
		return 0
	}
}

func sanitizeFilename(value string) string {
	trimmed := strings.TrimSpace(strings.ToLower(value))
	if trimmed == "" {
		return "certificate"
	}

	var builder strings.Builder
	lastDash := false
	for _, r := range trimmed {
		isAlphaNum := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9')
		if isAlphaNum {
			builder.WriteRune(r)
			lastDash = false
			continue
		}
		if !lastDash {
			builder.WriteRune('-')
			lastDash = true
		}
	}

	result := strings.Trim(builder.String(), "-")
	if result == "" {
		return "certificate"
	}
	return result
}

func publicKeySize(key any) int {
	switch typed := key.(type) {
	case *rsa.PublicKey:
		return typed.N.BitLen()
	default:
		return 0
	}
}

func getenv(key, fallback string) string {
	value := strings.TrimSpace(os.Getenv(key))
	if value == "" {
		return fallback
	}
	return value
}