// 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"
	"io/ioutil"
	"os"
	"path/filepath"
	pmpb "policy-manager/protos"
	"reflect"

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

const (
	// policyType is a Chrome OS constant to specify that the policy data
	// refers to a device policy.
	// See line 2 in third_party/chromium/devicepolicyproto/device_management_backend.proto.
	policyType = "google/chromeos/device"

	// configFilePerm is the permission for the COS policy file.
	configFilePerm = os.FileMode(0644)

	// 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
)

// 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(publicKeyFile string, privateKeyFile string) (*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 = ioutil.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 = ioutil.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(publicKeyFile string, privateKeyFile string) (*rsa.PrivateKey, error) {
	// Get private key.
	privateKeyBytes, err := ioutil.ReadFile(privateKeyFile)
	if err != nil {
		return nil, err
	}
	privateKey, err := x509.ParsePKCS1PrivateKey(privateKeyBytes)
	if err != nil {
		return nil, err
	}

	// Get public key.
	publicKeyBytes, err := ioutil.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(publicKeyFile string, privateKeyFile string, 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(publicKeyFile, privateKeyFile)
		}
		return nil, errors.New("missing key files")
	}

	// Check if valid public/private key files are present.
	_, errPrivateMissing := os.Stat(privateKeyFile)
	_, errPublicMissing := os.Stat(publicKeyFile)

	if os.IsNotExist(errPrivateMissing) || os.IsNotExist(errPublicMissing) {
		glog.V(1).Info("keys missing")
		return generateKeysIfNeeded()
	} else if privateKey, err := checkKeyFiles(publicKeyFile, privateKeyFile); err != nil {
		// The keys on disk do not form a valid private/public key pair.
		glog.V(1).Infof("invalid key pair: %v", err)
		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),
		UpdateStrategy: proto.String(""),
		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
}

// 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. 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 InitDevicePolicy(baseConfig *pmpb.InstanceConfig, devicePolicyFile string, instanceConfigFile string, publicKeyFile string, privateKeyFile string) error {
	// Make sure the policy directory exists.
	dir := filepath.Dir(instanceConfigFile)
	if err := os.MkdirAll(dir, configFilePerm); err != nil {
		return err
	}

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

	onDiskConfig, err := GetInstanceConfig(instanceConfigFile)
	// 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 SetInstanceConfig(onDiskConfig, devicePolicyFile, instanceConfigFile, publicKeyFile, privateKeyFile)
	}

	// Overwrite the copy of device policy on disk using baseConfig.
	return SetInstanceConfig(baseConfig, devicePolicyFile, instanceConfigFile, publicKeyFile, privateKeyFile)
}

// SetInstanceConfig sets a new device policy from an InstanceConfig. The
// new policy will atomically replace the existing policy.
func SetInstanceConfig(baseConfig *pmpb.InstanceConfig, devicePolicyFile string, instanceConfigFile string, publicKeyFile string, privateKeyFile string) 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(publicKeyFile, privateKeyFile, false)
	if err != nil {
		return err
	}

	// Read the current device policy. Ignore error if policy file is missing.
	policy := new(DevicePolicy)
	if err := policy.ReadPolicyFromDisk(instanceConfigFile); 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 cos device policy. Instance config changed from {%+v} to {%+v}",
			currentConfig, newConfig)
	}

	return policy.SignAndWrite(privateKey, devicePolicyFile, instanceConfigFile)
}

func GetInstanceConfig(instanceConfigFile string) (config *pmpb.InstanceConfig, err error) {
	// Read the current device policy.
	policy := new(DevicePolicy)
	err = policy.ReadPolicyFromDisk(instanceConfigFile)
	if err != nil {
		return nil, err
	}

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

	return config, nil
}