package server

import (
	"context"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"path/filepath"
	"strings"

	"github.com/pquerna/otp/totp"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/tstranex/u2f"
	"gopkg.in/yaml.v2"

	"git.autistici.org/ai3/go-common/clientutil"
	"git.autistici.org/ai3/go-common/pwhash"
	"git.autistici.org/id/auth"
)

// User contains the attributes of a user account as relevant to the
// authentication server. It is only used internally, to communicate
// between the authserver and its storage backends.
type User struct {
	Name                 string
	Email                string
	Shard                string
	EncryptedPassword    []byte
	TOTPSecret           string
	U2FRegistrations     []u2f.Registration
	AppSpecificPasswords []*AppSpecificPassword
	Groups               []string
}

// AppSpecificPassword is a password tied to a single service.
type AppSpecificPassword struct {
	Service           string
	EncryptedPassword []byte
}

// Has2FA returns true if the user supports any 2FA method.
func (u *User) Has2FA() bool {
	return u.HasU2F() || u.HasOTP()
}

// HasOTP returns true if the user supports (T)OTP.
func (u *User) HasOTP() bool {
	return u.TOTPSecret != ""
}

// HasU2F returns true if the user supports U2F.
func (u *User) HasU2F() bool {
	return len(u.U2FRegistrations) > 0
}

// UserInfo returns extra user information in the format required by
// the auth wire protocol.
func (u *User) UserInfo() *auth.UserInfo {
	return &auth.UserInfo{
		Email:  u.Email,
		Shard:  u.Shard,
		Groups: u.Groups,
	}
}

// UserBackend provides us with per-service user information.
type UserBackend interface {
	Close()
	GetUser(context.Context, *BackendSpec, string) (*User, bool)
}

// U2FShortTermStorage stores short-term u2f challenges.
type U2FShortTermStorage interface {
	SetUserChallenge(string, *u2f.Challenge) error
	GetUserChallenge(string) (*u2f.Challenge, bool)
}

type ratelimitKeyFunc func(*User, *auth.Request) string

func usernameKey(user *User, _ *auth.Request) string {
	return user.Name
}

func ipAddrKey(_ *User, req *auth.Request) string {
	if req.DeviceInfo != nil {
		return req.DeviceInfo.RemoteAddr
	}
	return ""
}

type authRatelimiterConfig struct {
	Limit         int      `yaml:"limit"`
	Period        int      `yaml:"period"`
	BlacklistTime int      `yaml:"blacklist_for"`
	OnFailure     bool     `yaml:"on_failure"`
	Keys          []string `yaml:"keys"`

	keyFuncs []ratelimitKeyFunc
}

func (r *authRatelimiterConfig) compile() error {
	for _, k := range r.Keys {
		var f ratelimitKeyFunc
		switch k {
		case "ip":
			f = ipAddrKey
		case "user":
			f = usernameKey
		default:
			return fmt.Errorf("unknown key %s", k)
		}
		r.keyFuncs = append(r.keyFuncs, f)
	}
	return nil
}

const rlKeySep = ";"

func (r *authRatelimiterConfig) key(user *User, req *auth.Request) string {
	if len(r.keyFuncs) == 1 {
		return r.keyFuncs[0](user, req)
	}

	var parts []string
	for _, f := range r.keyFuncs {
		parts = append(parts, f(user, req))
	}
	return strings.Join(parts, rlKeySep)
}

type authRatelimiter struct {
	*authRatelimiterConfig
	rl *Ratelimiter
}

func (r *authRatelimiter) AllowIncr(user *User, req *auth.Request) bool {
	return r.rl.AllowIncr(r.key(user, req))
}

type authBlacklist struct {
	*authRatelimiterConfig
	bl *Blacklist
}

func (b *authBlacklist) Allow(user *User, req *auth.Request) bool {
	return b.bl.Allow(b.key(user, req))
}

func (b *authBlacklist) Incr(user *User, req *auth.Request, resp *auth.Response) {
	if b.OnFailure && resp.Status == auth.StatusOK {
		return
	}
	b.bl.Incr(b.key(user, req))
}

type requestFilter interface {
	Filter(*User, *auth.Request, *auth.Response) *auth.Response
}

// BackendSpec specifies backend-specific configuration for a service.
type BackendSpec struct {
	LDAPSpec *LDAPServiceConfig `yaml:"ldap"`
	FileSpec string             `yaml:"file"`
}

// ServiceConfig defines the authentication backends for a service.
type ServiceConfig struct {
	BackendSpecs         []*BackendSpec `yaml:"backends"`
	ChallengeResponse    bool           `yaml:"challenge_response"`
	Enforce2FA           bool           `yaml:"enforce_2fa"`
	EnableDeviceTracking bool           `yaml:"enable_device_tracking"`

	Ratelimits []string `yaml:"rate_limits"`

	rl      []*authRatelimiter
	bl      []*authBlacklist
	filters []requestFilter
}

func (c *ServiceConfig) checkRateLimits(user *User, req *auth.Request) bool {
	for _, rl := range c.rl {
		if !rl.AllowIncr(user, req) {
			return false
		}
	}
	for _, bl := range c.bl {
		if !bl.Allow(user, req) {
			return false
		}
	}
	return true
}

func (c *ServiceConfig) notifyBlacklists(user *User, req *auth.Request, resp *auth.Response) {
	for _, bl := range c.bl {
		bl.Incr(user, req, resp)
	}
}

// Config for the authentication server.
type Config struct {
	// Global configuration for backends.
	LDAPConfig *LDAPConfig `yaml:"ldap_config"`

	// List of enabled backends.
	EnabledBackends []string `yaml:"enabled_backends"`

	// Service-specific configuration.
	Services map[string]*ServiceConfig `yaml:"services"`

	// Named rate limiter configurations.
	RateLimiters map[string]*authRatelimiterConfig `yaml:"rate_limits"`

	// Configuration for the user-meta-server backend.
	UserMetaDBConfig *clientutil.BackendConfig `yaml:"user_meta_server"`

	// Runtime versions of the above. These objects are shared by
	// all services, as they contain the actual map data.
	rl map[string]*Ratelimiter
	bl map[string]*Blacklist

	path string
}

func (c *Config) walkBackendSpecs(f func(*BackendSpec) error) error {
	for _, svc := range c.Services {
		for _, spec := range svc.BackendSpecs {
			if err := f(spec); err != nil {
				return err
			}
		}
	}
	return nil
}

func (c *Config) relativePath(path string) string {
	if strings.HasPrefix(path, "/") {
		return path
	}
	return filepath.Join(filepath.Dir(c.path), path)
}

func (c *Config) compile() error {
	// Build the global rate limiters and blacklists.
	c.rl = make(map[string]*Ratelimiter)
	c.bl = make(map[string]*Blacklist)
	for name, params := range c.RateLimiters {
		if err := params.compile(); err != nil {
			return err
		}
		if params.BlacklistTime > 0 {
			c.bl[name] = newBlacklist(params.Limit, params.Period, params.BlacklistTime)
		} else {
			c.rl[name] = newRatelimiter(params.Limit, params.Period)
		}
	}

	// Compile each service definition.
	for _, sc := range c.Services {
		for _, name := range sc.Ratelimits {
			config, ok := c.RateLimiters[name]
			if !ok {
				return fmt.Errorf("unknown rate limiter %s", name)
			}
			if rl, ok := c.rl[name]; ok {
				sc.rl = append(sc.rl, &authRatelimiter{config, rl})
			} else if bl, ok := c.bl[name]; ok {
				sc.bl = append(sc.bl, &authBlacklist{config, bl})
			} else {
				panic("can't find rl/bl")
			}
		}

		if sc.EnableDeviceTracking {
			if c.UserMetaDBConfig == nil {
				return errors.New("usermetadb config is missing")
			}
			dt, err := newDeviceFilter(c.UserMetaDBConfig)
			if err != nil {
				return err
			}
			sc.filters = append(sc.filters, dt)
		}
	}

	return nil
}

// LoadConfig loads the configuration from a YAML-encoded file.
func LoadConfig(path string) (*Config, error) {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, err
	}
	config := Config{path: path}
	if err := yaml.Unmarshal(data, &config); err != nil {
		return nil, err
	}
	if err := config.compile(); err != nil {
		return nil, err
	}
	return &config, nil
}

// Instrumentation.
var (
	authRequestsCounter = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "auth_requests",
			Help: "Number of authentication requests.",
		},
		[]string{"service", "status"},
	)
	ratelimitCounter = prometheus.NewCounterVec(
		prometheus.CounterOpts{
			Name: "auth_requests_ratelimited",
			Help: "Number of rate-limited authentication requests.",
		},
		[]string{"service"},
	)
)

func init() {
	prometheus.MustRegister(authRequestsCounter)
	prometheus.MustRegister(ratelimitCounter)
}

// Server is the main authentication server object.
type Server struct {
	backends     []UserBackend
	config       *Config
	u2fShortTerm U2FShortTermStorage
}

func newError() *auth.Response {
	return &auth.Response{Status: auth.StatusError}
}

func newOK() *auth.Response {
	return &auth.Response{Status: auth.StatusOK}
}

// NewServer creates a Server using the given configuration.
func NewServer(config *Config, u2fShortTerm U2FShortTermStorage) (*Server, error) {
	s := &Server{
		config:       config,
		u2fShortTerm: u2fShortTerm,
	}
	for _, name := range config.EnabledBackends {
		var b UserBackend
		var err error
		switch name {
		case "file":
			b, err = newFileBackend(config)
		case "ldap":
			b, err = newLDAPBackend(config)
		default:
			err = fmt.Errorf("unknown backend %s", name)
		}
		if err != nil {
			return nil, err
		}
		s.backends = append(s.backends, b)
	}
	return s, nil
}

// Close the authentication server and release all associated resources.
func (s *Server) Close() {
	for _, b := range s.backends {
		b.Close()
	}
}

func (s *Server) getServiceConfig(service string) (*ServiceConfig, bool) {
	c, ok := s.config.Services[service]
	if !ok {
		c, ok = s.config.Services["default"]
	}
	return c, ok
}

func (s *Server) getUser(ctx context.Context, serviceConfig *ServiceConfig, username string) (*User, bool) {
	for _, spec := range serviceConfig.BackendSpecs {
		for _, b := range s.backends {
			if user, ok := b.GetUser(ctx, spec, username); ok {
				return user, true
			}
		}
	}
	return nil, false
}

// Authenticate a user with the parameters specified in the incoming AuthRequest.
func (s *Server) Authenticate(ctx context.Context, req *auth.Request) *auth.Response {
	serviceConfig, ok := s.getServiceConfig(req.Service)
	if !ok {
		log.Printf("unknown service %s", req.Service)
		return newError()
	}

	user, ok := s.getUser(ctx, serviceConfig, req.Username)
	if !ok {
		// User is unknown to all backends. Do not proceed
		// further, but log and increment stats counters.
		log.Printf("unknown user %s", req.Username)
		authRequestsCounter.With(prometheus.Labels{
			"service": req.Service,
			"status":  "unknown_user",
		})
		return newError()
	}

	// Apply rate limiting and blacklisting _before_ invoking the
	// authentication handlers, as they may be CPU intensive.
	if allowed := serviceConfig.checkRateLimits(user, req); !allowed {
		ratelimitCounter.With(prometheus.Labels{
			"service": req.Service,
		}).Inc()
		return newError()
	}

	resp, err := s.authenticateUser(req, serviceConfig, user)
	if err != nil {
		resp = newError()
		log.Printf("auth: user=%s service=%s status=%s error=%s", req.Username, req.Service, resp.Status.String(), err)
	} else {
		// Log the request and response.
		log.Printf("auth: user=%s service=%s status=%s", req.Username, req.Service, resp.Status.String())
	}

	// Notify blacklists of the result.
	serviceConfig.notifyBlacklists(user, req, resp)

	// Increment stats counters.
	authRequestsCounter.With(prometheus.Labels{
		"service": req.Service,
		"status":  resp.Status.String(),
	}).Inc()

	return resp
}

// Authenticate a user. Returning an error should result in an
// AuthResponse with StatusError.
func (s *Server) authenticateUser(req *auth.Request, serviceConfig *ServiceConfig, user *User) (resp *auth.Response, err error) {
	// Verify different credentials depending on whether the user
	// has 2FA enabled or not, and on whether the service itself
	// supports challenge-response authentication.
	if serviceConfig.Enforce2FA || user.Has2FA() {
		if serviceConfig.ChallengeResponse {
			resp, err = s.authenticateUserWith2FA(user, req)
		} else {
			resp, err = s.authenticateUserWithASP(user, req)
		}
	} else {
		resp, err = s.authenticateUserWithPassword(user, req)
	}
	if err != nil {
		return
	}

	// Process the response through filters (device info checks,
	// etc) that may or may not change the response itself.
	for _, f := range serviceConfig.filters {
		if resp.Status == auth.StatusError {
			break
		}
		resp = f.Filter(user, req, resp)
	}

	// If the response is successful, augment it with user information.
	if resp.Status == auth.StatusOK {
		resp.UserInfo = user.UserInfo()
	}
	return
}

func (s *Server) authenticateUserWithPassword(user *User, req *auth.Request) (*auth.Response, error) {
	// Ok we only need to check the password here.
	if checkPassword(req.Password, user.EncryptedPassword) {
		return newOK(), nil
	}
	return nil, errors.New("wrong password")
}

func (s *Server) authenticateUserWithASP(user *User, req *auth.Request) (*auth.Response, error) {
	for _, asp := range user.AppSpecificPasswords {
		if asp.Service == req.Service && checkPassword(req.Password, asp.EncryptedPassword) {
			return newOK(), nil
		}
	}
	return nil, errors.New("wrong app-specific password")
}

func (s *Server) authenticateUserWith2FA(user *User, req *auth.Request) (*auth.Response, error) {
	// First of all verify the password.
	if !checkPassword(req.Password, user.EncryptedPassword) {
		return nil, errors.New("wrong password")
	}

	// If the request contains one of the 2FA attributes, verify
	// it. But if it contains none, we return with
	// AuthStatusInsufficientCredentials and potentially a 2FA
	// hint (for U2F).
	switch {
	case req.U2FResponse != nil:
		if user.HasU2F() && s.checkU2F(user, req.U2FResponse) {
			return newOK(), nil
		}
		return nil, errors.New("bad U2F response")
	case req.OTP != "":
		if user.HasOTP() && checkOTP(req.OTP, user.TOTPSecret) {
			return newOK(), nil
		}
		return nil, errors.New("bad OTP")
	default:
		resp := &auth.Response{
			Status: auth.StatusInsufficientCredentials,
		}
		if req.U2FAppID != "" && user.HasU2F() {
			resp.TFAMethod = auth.TFAMethodU2F
			signReq, err := s.u2fSignRequest(user, req.U2FAppID)
			if err != nil {
				return nil, err
			}
			resp.U2FSignRequest = signReq
		} else if user.HasOTP() {
			resp.TFAMethod = auth.TFAMethodOTP
		}
		return resp, nil
	}
}

func (s *Server) u2fSignRequest(user *User, appID string) (*u2f.WebSignRequest, error) {
	challenge, err := u2f.NewChallenge(appID, []string{appID})
	if err != nil {
		return nil, err
	}

	if err := s.u2fShortTerm.SetUserChallenge(user.Name, challenge); err != nil {
		return nil, err
	}

	return challenge.SignRequest(user.U2FRegistrations), nil
}

func (s *Server) checkU2F(user *User, resp *u2f.SignResponse) bool {
	challenge, ok := s.u2fShortTerm.GetUserChallenge(user.Name)
	if !ok {
		return false
	}
	for _, reg := range user.U2FRegistrations {
		_, err := reg.Authenticate(*resp, *challenge, 0)
		if err == nil {
			return true
		}
	}
	return false
}

func checkPassword(password, hash []byte) bool {
	return pwhash.ComparePassword(string(hash), string(password))
}

func checkOTP(otp, secret string) bool {
	return totp.Validate(otp, secret)
}