package main

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"encoding/xml"
	"fmt"
	"io"
	"net"
	"os"
	"strings"
)

type xmppStreamsFeatures struct {
	XMLName    xml.Name `xml:"http://etherx.jabber.org/streams features"`
}

type xmppTlsProceed struct {
	XMLName xml.Name `xml:"urn:ietf:params:xml:ns:xmpp-tls proceed"`
}

func tlsDial(snet, saddr string) (*tls.Conn, error) {
	switch snet {
	case "tcp":
		conn, err := tls.Dial(snet, saddr, &tls.Config{InsecureSkipVerify: true})
		if err != nil {
			return nil, err
		}
		return conn, nil
	case "xmpp":
		host, _, err := net.SplitHostPort(saddr)
		connTCP, err := net.Dial("tcp", saddr)
		if err != nil {
			return nil, err
		}

		decoder := xml.NewDecoder(connTCP)

		// send <stream> start
		_, err = fmt.Fprintf(connTCP, "<stream:stream xmlns='jabber:client' xmlns:stream='http://etherx.jabber.org/streams' to='%s' version='1.0'>", host)
		if err != nil {
			return nil, err
		}
		// read <stream> start
		for {
			t, err := decoder.Token()
			if err != nil || t == nil {
				return nil, err
			}
			if se, ok := t.(xml.StartElement); ok {
				if se.Name.Local != "stream" {
					return nil, xml.UnmarshalError(fmt.Sprintf("expected element of type <%s> but have <%s>", "stream", se.Name.Local))
				}
				break
			}
		}
		// read <features>
		var features xmppStreamsFeatures
		err = decoder.DecodeElement(&features, nil)
		if err != nil {
			return nil, err
		}
		// send <starttls>
		_, err = io.WriteString(connTCP, "<starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'/>")
		if err != nil {
			return nil, err
		}
		// read <proceed>
		var proceed xmppTlsProceed
		err = decoder.DecodeElement(&proceed, nil)
		if err != nil {
			return nil, err
		}

		connTLS := tls.Client(connTCP, &tls.Config{InsecureSkipVerify: true})
		err = connTLS.Handshake()
		if err != nil {
			return nil, err
		}
		return connTLS, nil
	default:
		return nil, fmt.Errorf("Unknown TLS network: %q", snet)
	}
}

func getcert(socket string) (*x509.Certificate, error){
	snet, saddr := split(socket)
	host, _, err := net.SplitHostPort(saddr)
	if err != nil {
		return nil, err
	}
	conn, err := tlsDial(snet, saddr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	cstate := conn.ConnectionState()

	opts := x509.VerifyOptions{
		DNSName: host,
		Intermediates: x509.NewCertPool(),
	}
	for _, cert := range cstate.PeerCertificates[1:] {
		opts.Intermediates.AddCert(cert)
	}

	cert := cstate.PeerCertificates[0]
	_, err = cert.Verify(opts)
	return cert, err
}

func split(socket string) (net, addr string) {
	ary := strings.SplitN(socket, ":", 2)
	if len(ary) == 1 {
		return "tcp", ary[0]
	}
	return ary[0], ary[1]
}

func main() {
	for _, socket := range os.Args[1:] {
		cert, err := getcert(socket)
		if cert == nil {
			fmt.Fprintf(os.Stderr, "Could not get certificate for socket %q: %q\n", socket, err)
			os.Exit(1)
		}
		block := pem.Block{
			Type: "CERTIFICATE",
			Headers: map[string]string{"X-Socket": socket},
			Bytes: cert.Raw,
		}
		if err != nil {
			block.Headers["X-Error"] = err.Error()
		}
		pem.Encode(os.Stdout, &block)
	}
}