// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package devicepolicy

// Implementation of the Manager interface for managing Chrome OS device policy.
// For more detail on Chrome OS device policy, see
// http://www.chromium.org/developers/how-tos/enterprise/protobuf-encoded-policy-blobs

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"errors"
	"os"
	"path/filepath"
	"reflect"

	pmpb "policy_manager/policymanagerproto"
	"policy_manager/sysapi"

	"github.com/golang/glog"
	"github.com/golang/protobuf/proto"
)

const (
	// publicKeyFile is the file containing the public key that will be used
	// to verify the signature on the device policy file.
	// The key will be encoded using DER-encoded PKIX format.
	// Note that the file name is decided by Chrome OS.
	publicKeyFile = "/var/lib/whitelist/owner.key"

	// privateKeyFile is the file containing the private key that will be
	// used to sign the policy file.
	// The key will be encoded using ASN.1 DER encoded form.
	// Note that we are signing the policy file locally since we do not want
	// to do any key management on the server side.
	privateKeyFile = "/var/lib/whitelist/private.key"

	// keyFilePerm is the permission for public/private key files.
	keyFilePerm = os.FileMode(0600)

	// rsaKeyBits is the number of bits to use to generate the private key.
	rsaKeyBits = 1024
)

// managerImpl implements the Manager interface.
type managerImpl struct {
	// api is the interface used to call system API.
	api sysapi.APIHandler
}

// createKeyFiles creates a new public and private key pair and writes them to
// disk.
// Note that when the key files are being created, any existing policy cannot be
// read and verified by Chrome OS components such as update_engine.
func createKeyFiles(api sysapi.APIHandler) (*rsa.PrivateKey, error) {
	// Generate RSA key pair.
	key, err := rsa.GenerateKey(rand.Reader, rsaKeyBits)
	if err != nil {
		return nil, err
	}

	// Serialize the public key to DER-encoded PKIX format and output to
	// file.
	publicKeyBytes, err := x509.MarshalPKIXPublicKey(key.Public())
	if err != nil {
		return nil, err
	}
	err = api.WriteFile(publicKeyFile, publicKeyBytes, keyFilePerm)
	if err != nil {
		return nil, err
	}

	// Serialize the private key to ASN.1 DER encoded form and output to
	// file.
	privateKeyBytes := x509.MarshalPKCS1PrivateKey(key)
	if err != nil {
		return nil, err
	}
	err = api.WriteFile(privateKeyFile, privateKeyBytes, keyFilePerm)
	if err != nil {
		return nil, err
	}

	glog.V(1).Info("new keys created")
	return key, nil
}

// checkKeyFiles checks the private/public key files are present on disk and
// form a valid key pair.
// An error is returned if the private or public key file is not present or if
// the keys don't match.
// If no error, then the private key is returned.
func checkKeyFiles(api sysapi.APIHandler) (*rsa.PrivateKey, error) {
	// Get private key.
	privateKeyBytes, err := api.ReadFile(privateKeyFile)
	if err != nil {
		return nil, err
	}
	privateKey, err := x509.ParsePKCS1PrivateKey(privateKeyBytes)
	if err != nil {
		return nil, err
	}

	// Get public key.
	publicKeyBytes, err := api.ReadFile(publicKeyFile)
	if err != nil {
		return nil, err
	}
	publicKey, err := x509.ParsePKIXPublicKey(publicKeyBytes)
	if err != nil {
		return nil, err
	}

	// Verify that the public and private keys match.
	if !reflect.DeepEqual(publicKey, privateKey.Public()) {
		return nil, errors.New("public and private keys don't match")
	}

	return privateKey, nil
}

// getKeys returns the private key of the public/private key pair present on
// disk. If no valid key pair exsits, an error will be returned unless the
// caller has set generateIfMissing to true.
func getKeys(api sysapi.APIHandler, generateIfMissing bool) (*rsa.PrivateKey, error) {
	// generateKeysIfNeeded is called when there is no valid key pair on
	// disk and should consult the generateIfMissing flag to decide wheher
	// new keys should be generated.
	generateKeysIfNeeded := func() (*rsa.PrivateKey, error) {
		if generateIfMissing {
			return createKeyFiles(api)
		}
		return nil, errors.New("missing key files")
	}

	// Check if valid public/private key files are present.
	if !api.FileExists(privateKeyFile) || !api.FileExists(publicKeyFile) {
		glog.V(1).Info("keys missing")
		return generateKeysIfNeeded()
	} else if privateKey, err := checkKeyFiles(api); err != nil {
		// The keys on disk do not form a valid private/public key pair.
		glog.V(1).Info("invalid key pair")
		return generateKeysIfNeeded()
	} else {
		return privateKey, nil
	}
}

// getDefaultInstanceConfigFromBase returns the instance config derived from given
// baseConfig. All missing fields in baseConfig are populated with default values.
func getDefaultInstanceConfigFromBase(baseConfig *pmpb.InstanceConfig) *pmpb.InstanceConfig {
	// default policy allows all updates, disables cos monitoring logging and monitoring.
	result := &pmpb.InstanceConfig{
		MetricsEnabled:       proto.Bool(false),
		TargetVersionPrefix:  proto.String(""),
		RebootAfterUpdate:    proto.Bool(false),
		UpdateScatterSeconds: proto.Int64(33),
		HealthMonitorConfig: &pmpb.HealthMonitorConfig{
			Enforced:          proto.Bool(false),
			LoggingEnabled:    proto.Bool(false),
			MonitoringEnabled: proto.Bool(false),
		},
	}

	// Override fields from baseConfig.
	proto.Merge(result, baseConfig)
	return result
}

// NewManager returns a new Manager.
func NewManager(api sysapi.APIHandler) Manager {
	return &managerImpl{api}
}

// InitDevicePolicy initializes instance's environment to support device policy.
// The initialization process includes (in order):
// 1. Create the directory to store policy file if it doesn't exist.
// 2. Detect if a valid public/private key pair is present. If not, create a new
//      key pair and store them on disk. In addition, the private key is stored
//      in memory with the manager.
// 3. New devicepolicy will be created and written to disk (overwriting any previous
// settings), All missing fields will have default values. If the given InstanceConfig
// is valid (not nil), new devicepolicy will be generated from it; if the given
// InstanceConfig is nil, which indicates metadata server fetching error, new devicepolicy
// will be generated from current devicepolicy, which is stored on disk.
func (manager *managerImpl) InitDevicePolicy(baseConfig *pmpb.InstanceConfig) error {
	// Make sure the policy directory exists.
	dir := filepath.Dir(policyFile)
	if err := manager.api.MkdirAll(dir, policyFilePerm); err != nil {
		return err
	}

	// Create new keys if no valid keys are present.
	_, err := getKeys(manager.api, true)
	if err != nil {
		return err
	}

	onDiskConfig, err := manager.GetInstanceConfig()
	// If failed to fetch config from metadata server, use the copy on disk.
	if baseConfig == nil && err == nil {
		glog.V(1).Infof("Reusing existing device policy for initialization: %v", onDiskConfig)
		return manager.SetInstanceConfig(onDiskConfig)
	}

	// Overwrite the copy of device policy on disk using baseConfig.
	return manager.SetInstanceConfig(baseConfig)
}

// SetInstanceConfig sets a new device policy from an InstanceConfig. The
// new policy will atomically replace the existing policy.
// This function requires that a valid public key is present before proceeding
// with generating the policy file. This is to prevent the situation where the
// policy file is written but not the public key file, causing Chrome OS
// components to ignore the policy because the signature cannot be verified.
func (manager *managerImpl) SetInstanceConfig(baseConfig *pmpb.InstanceConfig) error {
	// We need to populate unspecified fields with default values.
	config := getDefaultInstanceConfigFromBase(baseConfig)

	// Get the private key from disk. Don't generate new keys if valid
	// public/private key pair is missing.
	privateKey, err := getKeys(manager.api, false)
	if err != nil {
		return err
	}
	publicKey := privateKey.Public().(*rsa.PublicKey)

	// Read the current device policy. Ignore error if policy file is
	// missing.
	policy := new(DevicePolicy)
	err = policy.ReadPolicyFromDisk(publicKey, manager.api, true)
	if err != nil {
		return err
	}

	currentConfig := new(pmpb.InstanceConfig)
	policy.ToInstanceConfig(currentConfig)

	// Set the device policy based on new config.
	policyChanged := policy.SetFromInstanceConfig(config)
	if !policyChanged {
		glog.Infof("Skipping device policy update. Current: {%+v}", currentConfig)
		return nil
	} else {
		newConfig := new(pmpb.InstanceConfig)
		policy.ToInstanceConfig(newConfig)
		glog.Infof("Updating device policy. Instance config changed from {%+v} to {%+v}",
			currentConfig, newConfig)
	}

	err = policy.SignAndWrite(privateKey, manager.api)

	if err == nil {
		dbus_error := signalPropertyChange()
		if dbus_error != nil {
			glog.Warningf("Can't send PropertyChange signal on D-bus. "+
				"Some consumers may not get device policy update: %s", dbus_error)
		}
	}
	return err
}

func (manager *managerImpl) GetInstanceConfig() (config *pmpb.InstanceConfig, err error) {
	// Get the private key from disk. Error out if they are missing.
	privateKey, err := getKeys(manager.api, false)
	if err != nil {
		return nil, err
	}
	publicKey := privateKey.Public().(*rsa.PublicKey)

	// Read the current device policy.
	policy := new(DevicePolicy)
	err = policy.ReadPolicyFromDisk(publicKey, manager.api, false)
	if err != nil {
		return nil, err
	}

	config = new(pmpb.InstanceConfig)
	policy.ToInstanceConfig(config)

	return config, nil
}