package tor

import (
	"bytes"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"io/ioutil"
	"net/textproto"
	"os"
	"strings"
	"sync/atomic"
)

const (
	// success is the Tor Control response code representing a successful
	// request.
	success = 250

	// nonceLen is the length of a nonce generated by either the controller
	// or the Tor server
	nonceLen = 32

	// cookieLen is the length of the authentication cookie.
	cookieLen = 32

	// ProtocolInfoVersion is the `protocolinfo` version currently supported
	// by the Tor server.
	ProtocolInfoVersion = 1
)

var (
	// serverKey is the key used when computing the HMAC-SHA256 of a message
	// from the server.
	serverKey = []byte("Tor safe cookie authentication " +
		"server-to-controller hash")

	// controllerKey is the key used when computing the HMAC-SHA256 of a
	// message from the controller.
	controllerKey = []byte("Tor safe cookie authentication " +
		"controller-to-server hash")
)

// Controller is an implementation of the Tor Control protocol. This is used in
// order to communicate with a Tor server. Its only supported method of
// authentication is the SAFECOOKIE method.
//
// NOTE: The connection to the Tor server must be authenticated before
// proceeding to send commands. Otherwise, the connection will be closed.
//
// TODO:
//   * if adding support for more commands, extend this with a command queue?
//   * place under sub-package?
//   * support async replies from the server
type Controller struct {
	// started is used atomically in order to prevent multiple calls to
	// Start.
	started int32

	// stopped is used atomically in order to prevent multiple calls to
	// Stop.
	stopped int32

	// conn is the underlying connection between the controller and the
	// Tor server. It provides read and write methods to simplify the
	// text-based messages within the connection.
	conn *textproto.Conn

	// controlAddr is the host:port the Tor server is listening locally for
	// controller connections on.
	controlAddr string
}

// NewController returns a new Tor controller that will be able to interact with
// a Tor server.
func NewController(controlAddr string) *Controller {
	return &Controller{controlAddr: controlAddr}
}

// Start establishes and authenticates the connection between the controller and
// a Tor server. Once done, the controller will be able to send commands and
// expect responses.
func (c *Controller) Start() error {
	if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
		return nil
	}

	conn, err := textproto.Dial("tcp", c.controlAddr)
	if err != nil {
		return fmt.Errorf("unable to connect to Tor server: %v", err)
	}

	c.conn = conn

	return c.authenticate()
}

// Stop closes the connection between the controller and the Tor server.
func (c *Controller) Stop() error {
	if !atomic.CompareAndSwapInt32(&c.stopped, 0, 1) {
		return nil
	}

	return c.conn.Close()
}

// sendCommand sends a command to the Tor server and returns its response, as a
// single space-delimited string, and code.
func (c *Controller) sendCommand(command string) (int, string, error) {
	if err := c.conn.Writer.PrintfLine(command); err != nil {
		return 0, "", err
	}

	// We'll use ReadResponse as it has built-in support for multi-line
	// text protocol responses.
	code, reply, err := c.conn.Reader.ReadResponse(success)
	if err != nil {
		return code, reply, err
	}

	return code, reply, nil
}

// parseTorReply parses the reply from the Tor server after receving a command
// from a controller. This will parse the relevent reply parameters into a map
// of keys and values.
func parseTorReply(reply string) map[string]string {
	params := make(map[string]string)

	// Replies can either span single or multiple lines, so we'll default
	// to stripping whitespace and newlines in order to retrieve the
	// individual contents of it. The -1 indicates that we want this to span
	// across all instances of a newline.
	contents := strings.Split(strings.Replace(reply, "\n", " ", -1), " ")
	for _, content := range contents {
		// Each parameter within the reply should be of the form
		// "KEY=VALUE". If the parameter doesn't contain "=", then we
		// can assume it does not provide any other relevant information
		// already known.
		keyValue := strings.Split(content, "=")
		if len(keyValue) != 2 {
			continue
		}

		key := keyValue[0]
		value := keyValue[1]
		params[key] = value
	}

	return params
}

// authenticate authenticates the connection between the controller and the
// Tor server using the SAFECOOKIE authentication method.
func (c *Controller) authenticate() error {
	// Before proceeding to authenticate the connection, we'll retrieve
	// the authentication cookie of the Tor server. This will be used
	// throughout the authentication routine. We do this before as once the
	// authentication routine has begun, it is not possible to retrieve it
	// mid-way.
	cookie, err := c.getAuthCookie()
	if err != nil {
		return fmt.Errorf("unable to retrieve authentication cookie: "+
			"%v", err)
	}

	// Authenticating using the SAFECOOKIE authentication method is a two
	// step process. We'll kick off the authentication routine by sending
	// the AUTHCHALLENGE command followed by a hex-encoded 32-byte nonce.
	clientNonce := make([]byte, nonceLen)
	if _, err := rand.Read(clientNonce); err != nil {
		return fmt.Errorf("unable to generate client nonce: %v", err)
	}

	cmd := fmt.Sprintf("AUTHCHALLENGE SAFECOOKIE %x", clientNonce)
	_, reply, err := c.sendCommand(cmd)
	if err != nil {
		return err
	}

	// If successful, the reply from the server should be of the following
	// format:
	//
	//	"250 AUTHCHALLENGE"
	//		SP "SERVERHASH=" ServerHash
	//		SP "SERVERNONCE=" ServerNonce
	//		CRLF
	//
	// We're interested in retrieving the SERVERHASH and SERVERNONCE
	// parameters, so we'll parse our reply to do so.
	replyParams := parseTorReply(reply)

	// Once retrieved, we'll ensure these values are of proper length when
	// decoded.
	serverHash, ok := replyParams["SERVERHASH"]
	if !ok {
		return errors.New("server hash not found in reply")
	}
	decodedServerHash, err := hex.DecodeString(serverHash)
	if err != nil {
		return fmt.Errorf("unable to decode server hash: %v", err)
	}
	if len(decodedServerHash) != sha256.Size {
		return errors.New("invalid server hash length")
	}

	serverNonce, ok := replyParams["SERVERNONCE"]
	if !ok {
		return errors.New("server nonce not found in reply")
	}
	decodedServerNonce, err := hex.DecodeString(serverNonce)
	if err != nil {
		return fmt.Errorf("unable to decode server nonce: %v", err)
	}
	if len(decodedServerNonce) != nonceLen {
		return errors.New("invalid server nonce length")
	}

	// The server hash above was constructed by computing the HMAC-SHA256
	// of the message composed of the cookie, client nonce, and server
	// nonce. We'll redo this computation ourselves to ensure the integrity
	// and authentication of the message.
	hmacMessage := bytes.Join(
		[][]byte{cookie, clientNonce, decodedServerNonce}, []byte{},
	)
	computedServerHash := computeHMAC256(serverKey, hmacMessage)
	if !hmac.Equal(computedServerHash, decodedServerHash) {
		return fmt.Errorf("expected server hash %x, got %x",
			decodedServerHash, computedServerHash)
	}

	// If the MAC check was successful, we'll proceed with the last step of
	// the authentication routine. We'll now send the AUTHENTICATE command
	// followed by a hex-encoded client hash constructed by computing the
	// HMAC-SHA256 of the same message, but this time using the controller's
	// key.
	clientHash := computeHMAC256(controllerKey, hmacMessage)
	if len(clientHash) != sha256.Size {
		return errors.New("invalid client hash length")
	}

	cmd = fmt.Sprintf("AUTHENTICATE %x", clientHash)
	if _, _, err := c.sendCommand(cmd); err != nil {
		return err
	}

	return nil
}

// getAuthCookie retrieves the authentication cookie in bytes from the Tor
// server. Cookie authentication must be enabled for this to work.
func (c *Controller) getAuthCookie() ([]byte, error) {
	// Retrieve the authentication methods currently supported by the Tor
	// server.
	authMethods, cookieFilePath, _, err := c.ProtocolInfo()
	if err != nil {
		return nil, err
	}

	// Ensure that the Tor server supports the SAFECOOKIE authentication
	// method.
	safeCookieSupport := false
	for _, authMethod := range authMethods {
		if authMethod == "SAFECOOKIE" {
			safeCookieSupport = true
		}
	}

	if !safeCookieSupport {
		return nil, errors.New("the Tor server is currently not " +
			"configured for cookie authentication")
	}

	// Read the cookie from the file and ensure it has the correct length.
	cookie, err := ioutil.ReadFile(cookieFilePath)
	if err != nil {
		return nil, err
	}

	if len(cookie) != cookieLen {
		return nil, errors.New("invalid authentication cookie length")
	}

	return cookie, nil
}

// computeHMAC256 computes the HMAC-SHA256 of a key and message.
func computeHMAC256(key, message []byte) []byte {
	mac := hmac.New(sha256.New, key)
	mac.Write(message)
	return mac.Sum(nil)
}

// ProtocolInfo returns the different authentication methods supported by the
// Tor server and the version of the Tor server.
func (c *Controller) ProtocolInfo() ([]string, string, string, error) {
	// We'll start off by sending the "PROTOCOLINFO" command to the Tor
	// server. We should receive a reply of the following format:
	//
	//	METHODS=COOKIE,SAFECOOKIE
	//	COOKIEFILE="/home/user/.tor/control_auth_cookie"
	//	VERSION Tor="0.3.2.10"
	//
	// We're interested in retrieving all of these fields, so we'll parse
	// our reply to do so.
	cmd := fmt.Sprintf("PROTOCOLINFO %d", ProtocolInfoVersion)
	_, reply, err := c.sendCommand(cmd)
	if err != nil {
		return nil, "", "", err
	}

	info := parseTorReply(reply)
	methods, ok := info["METHODS"]
	if !ok {
		return nil, "", "", errors.New("auth methods not found in " +
			"reply")
	}

	cookieFile, ok := info["COOKIEFILE"]
	if !ok {
		return nil, "", "", errors.New("cookie file path not found " +
			"in reply")
	}

	version, ok := info["Tor"]
	if !ok {
		return nil, "", "", errors.New("Tor version not found in reply")
	}

	// Finally, we'll clean up the results before returning them.
	authMethods := strings.Split(methods, ",")
	cookieFilePath := strings.Trim(cookieFile, "\"")
	torVersion := strings.Trim(version, "\"")

	return authMethods, cookieFilePath, torVersion, nil
}

// VirtToTargPorts is a mapping of virtual ports to target ports. When creating
// an onion service, it will be listening externally on each virtual port. Each
// virtual port can then be mapped to one or many target ports internally. When
// accessing the onion service at a specific virtual port, it will forward the
// traffic to a mapped randomly chosen target port.
type VirtToTargPorts = map[int]map[int]struct{}

// AddOnionV2 creates a new v2 onion service and returns its onion address(es).
// Once created, the new onion service will remain active until the connection
// between the controller and the Tor server is closed. The path to a private
// key can be provided in order to restore a previously created onion service.
// If a file at this path does not exist, a new onion service will be created
// and its private key will be saved to a file at this path. A mapping of
// virtual ports to target ports should also be provided. Each virtual port will
// be the ports where the onion service can be reached at, while the mapped
// target ports will be the ports where the onion service is running locally.
func (c *Controller) AddOnionV2(privateKeyFilename string,
	virtToTargPorts VirtToTargPorts) ([]*OnionAddr, error) {

	// We'll start off by checking if the file containing the private key
	// exists. If it does not, then we should request the server to create
	// a new onion service and return its private key. Otherwise, we'll
	// request the server to recreate the onion server from our private key.
	var keyParam string
	if _, err := os.Stat(privateKeyFilename); os.IsNotExist(err) {
		keyParam = "NEW:RSA1024"
	} else {
		privateKey, err := ioutil.ReadFile(privateKeyFilename)
		if err != nil {
			return nil, err
		}
		keyParam = string(privateKey)
	}

	// Now, we'll determine the different virtual ports on which this onion
	// service will be accessed by.
	var portParam string
	for virtPort, targPorts := range virtToTargPorts {
		// If the virtual port doesn't map to any target ports, we'll
		// use the virtual port as the target port.
		if len(targPorts) == 0 {
			portParam += fmt.Sprintf("Port=%d,%d ", virtPort,
				virtPort)
			continue
		}

		// Otherwise, we'll create a mapping from the virtual port to
		// each target port.
		for targPort := range targPorts {
			portParam += fmt.Sprintf("Port=%d,%d ", virtPort,
				targPort)
		}
	}

	cmd := fmt.Sprintf("ADD_ONION %s %s", keyParam, portParam)
	_, reply, err := c.sendCommand(cmd)
	if err != nil {
		return nil, err
	}

	// If successful, the reply from the server should be of the following
	// format, depending on whether a private key has been requested:
	//
	//	C: ADD_ONION RSA1024:[Blob Redacted] Port=80,8080
	//	S: 250-ServiceID=testonion1234567
	//	S: 250 OK
	//
	//	C: ADD_ONION NEW:RSA1024 Port=80,8080
	//	S: 250-ServiceID=testonion1234567
	//	S: 250-PrivateKey=RSA1024:[Blob Redacted]
	//	S: 250 OK
	//
	// We're interested in retrieving the service ID, which is the public
	// name of the service, and the private key if requested.
	replyParams := parseTorReply(reply)
	serviceID, ok := replyParams["ServiceID"]
	if !ok {
		return nil, errors.New("service id not found in reply")
	}

	// If a new onion service was created, we'll write its private key to
	// disk under strict permissions in the event that it needs to be
	// recreated later on.
	if privateKey, ok := replyParams["PrivateKey"]; ok {
		err := ioutil.WriteFile(
			privateKeyFilename, []byte(privateKey), 0600,
		)
		if err != nil {
			return nil, fmt.Errorf("unable to write private key "+
				"to file: %v", err)
		}
	}

	// Finally, return the different onion addresses composed of the service
	// ID, along with the onion suffix, and the different virtual ports this
	// onion service can be reached at.
	onionService := serviceID + ".onion"
	addrs := make([]*OnionAddr, 0, len(virtToTargPorts))
	for virtPort := range virtToTargPorts {
		addr := &OnionAddr{
			OnionService: onionService,
			Port:         virtPort,
		}
		addrs = append(addrs, addr)
	}

	return addrs, nil
}