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

// Utility functions for manipulating device policy.

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"errors"
	"os"

	pmpb "policy_manager/policymanagerproto"
	"policy_manager/sysapi"
	pb "policy_manager/third_party/chromium/devicepolicyproto"

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

const (
	// policyFile is the file of the device policy.
	policyFile = "/var/lib/whitelist/policy"

	// policyFilePerm is the permission for policy related files.
	policyFilePerm = os.FileMode(0644)

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

	// cosDevicePolicyFile is the file storing the COS policy.
	cosDevicePolicyFile = "/var/lib/whitelist/cos_device_policy"

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

// hashPolicyData hashes the given policy data using SHA1.
func hashPolicyData(data []byte) []byte {
	h := sha1.New()
	h.Write(data)
	return h.Sum(nil)
}

// verifySignature verifies the device policy's signature using the provided
// public key.
func verifySignature(policy *pb.PolicyFetchResponse, key *rsa.PublicKey) error {
	if policy == nil {
		return errors.New("policy is nil")
	}

	// Hash policy data.
	hash := hashPolicyData(policy.GetPolicyData())

	// Verify the signature.
	err := rsa.VerifyPKCS1v15(key, crypto.SHA1, hash,
		policy.GetPolicyDataSignature())
	if err != nil {
		return err
	}

	return nil
}

// DevicePolicy is an opaque type that represents the device policy. This hides
// the complexity of the device policy protobufs from the users.
type DevicePolicy struct {
	// autoUpdateSettings is an AutoUpdateSettingsProto that controls the
	// behavior of update_engine.
	autoUpdateSettings *pb.AutoUpdateSettingsProto

	metricsEnabled *pb.MetricsEnabledProto

	cosDevicePolicy *pmpb.CosDevicePolicy
}

// generateDevicePolicy generates a signed PolicyFetchResponse protobuf that can
// be written to disk. The provided private key will be used to generate the
// signature.
func (p *DevicePolicy) generateDevicePolicy(key *rsa.PrivateKey,
	api sysapi.APIHandler) (*pb.PolicyFetchResponse, error) {
	// Generate ChromeDeviceSettingsProto.
	chromeDeviceSettings := pb.ChromeDeviceSettingsProto{
		AutoUpdateSettings: p.autoUpdateSettings,
		MetricsEnabled:     p.metricsEnabled,
	}
	chromeDeviceSettingsBytes, err := proto.Marshal(&chromeDeviceSettings)
	if err != nil {
		return nil, err
	}

	// Generate PolicyData.
	policyData := pb.PolicyData{
		PolicyType:       proto.String(policyType),
		Timestamp:        proto.Int64(api.GetUnixTimestamp()),
		PolicyValue:      chromeDeviceSettingsBytes,
		PublicKeyVersion: proto.Int32(1),
	}
	policyDataBytes, err := proto.Marshal(&policyData)
	if err != nil {
		return nil, err
	}

	// Hash policyData.
	hash := hashPolicyData(policyDataBytes)

	// Sign PolicyData.
	policyDataSig, err := rsa.SignPKCS1v15(rand.Reader, key, crypto.SHA1, hash)
	if err != nil {
		return nil, err
	}

	// Generate PolicyFetchResponse.
	policyFetchResponse := &pb.PolicyFetchResponse{
		PolicyData:          policyDataBytes,
		PolicyDataSignature: policyDataSig,
	}

	return policyFetchResponse, nil
}

// SetFromInstanceConfig sets the settings from the instance config to the
// DevicePolicy.
// Returns 'true' if this resulted in any change in device policy, otherwise false.
func (p *DevicePolicy) SetFromInstanceConfig(config *pmpb.InstanceConfig) bool {
	if config == nil {
		return false
	}
	policyChanged := false

	autoUpdateSettings := &pb.AutoUpdateSettingsProto{
		TargetVersionPrefix:    proto.String(config.GetTargetVersionPrefix()),
		RebootAfterUpdate:      proto.Bool(config.GetRebootAfterUpdate()),
		ScatterFactorInSeconds: proto.Int64(config.GetUpdateScatterSeconds()),
	}
	metricsEnabled := &pb.MetricsEnabledProto{
		MetricsEnabled: proto.Bool(config.GetMetricsEnabled()),
	}

	cosDevicePolicy := &pmpb.CosDevicePolicy{
		HealthMonitorConfig: proto.Clone(config.HealthMonitorConfig).(*pmpb.HealthMonitorConfig),
	}

	if p.autoUpdateSettings == nil || !proto.Equal(autoUpdateSettings, p.autoUpdateSettings) {
		policyChanged = true
	}

	if p.metricsEnabled == nil || !proto.Equal(metricsEnabled, p.metricsEnabled) {
		policyChanged = true
	}

	if p.cosDevicePolicy == nil || !proto.Equal(cosDevicePolicy, p.cosDevicePolicy) {
		policyChanged = true
	}

	p.autoUpdateSettings = autoUpdateSettings
	p.metricsEnabled = metricsEnabled
	p.cosDevicePolicy = cosDevicePolicy
	return policyChanged
}

// ToInstanceConfig converts the DevicePolicy into InstanceConfig proto (i.e.
// reverse of the above function).
func (p *DevicePolicy) ToInstanceConfig(config *pmpb.InstanceConfig) {
	if config == nil {
		return
	}

	if p.autoUpdateSettings != nil {
		config.TargetVersionPrefix = proto.String(p.autoUpdateSettings.GetTargetVersionPrefix())
		config.RebootAfterUpdate = proto.Bool(p.autoUpdateSettings.GetRebootAfterUpdate())
		config.UpdateScatterSeconds = proto.Int64(p.autoUpdateSettings.GetScatterFactorInSeconds())
	}

	if p.metricsEnabled != nil {
		config.MetricsEnabled = proto.Bool(p.metricsEnabled.GetMetricsEnabled())
	}

	if p.cosDevicePolicy != nil {
		config.HealthMonitorConfig = proto.Clone(p.cosDevicePolicy.HealthMonitorConfig).(*pmpb.HealthMonitorConfig)
	}
}

// SignAndWrite signs the device policy using the given key and writes the
// policy to disk. The policy will atomically replace the existing policy.
func (p *DevicePolicy) SignAndWrite(key *rsa.PrivateKey, api sysapi.APIHandler) error {
	// Generate policy protobuf.
	policyFetchResponse, err := p.generateDevicePolicy(key, api)
	if err != nil {
		return nil
	}

	policyBytes, err := proto.Marshal(policyFetchResponse)
	if err != nil {
		return err
	}
	cosDevicePolicyBytes, err := proto.Marshal(p.cosDevicePolicy)
	if err != nil {
		return err
	}

	// Write the policy file atomically.
	if err := api.AtomicWriteFile(policyFile, policyBytes, policyFilePerm); err != nil {
		return err
	}
	if err := api.AtomicWriteFile(cosDevicePolicyFile, cosDevicePolicyBytes, cosDevicePolicyFilePerm); err != nil {
		return err
	}

	return nil
}

// ReadPolicyFromDisk reads the device policy from disk and verifies its
// signature using the provided public key. An error is also returned if the
// policy file on disk does not contain a Chrome OS device policy.
// No error is returned for missing policy file if ignoreMissing is set to true.
func (p *DevicePolicy) ReadPolicyFromDisk(key *rsa.PublicKey,
	api sysapi.APIHandler, ignoreMissing bool) error {

	if !api.FileExists(policyFile) || !api.FileExists(cosDevicePolicyFile) {
		if ignoreMissing {
			return nil
		}
		return errors.New("missing policy file")
	}

	// Read the policy file and convert it to protobuf format.
	policyBytes, err := api.ReadFile(policyFile)
	if err != nil {
		return err
	}
	cosDevicePolicyBytes, err := api.ReadFile(cosDevicePolicyFile)
	if err != nil {
		return err
	}

	var policy pb.PolicyFetchResponse
	err = proto.Unmarshal(policyBytes, &policy)
	if err != nil {
		return err
	}

	// Verify the signature.
	if err := verifySignature(&policy, key); err != nil {
		return err
	}

	// Unmarshal PolicyData.
	var policyData pb.PolicyData
	err = proto.Unmarshal(policy.PolicyData, &policyData)
	if err != nil {
		return err
	}

	// Make sure policy type is Chrome OS device policy.
	if policyData.GetPolicyType() != policyType {
		return errors.New("wrong policy type")
	}

	// Unmarshal ChromeDeviceSettingsProto.
	var chromeDeviceSettings pb.ChromeDeviceSettingsProto
	err = proto.Unmarshal(policyData.PolicyValue, &chromeDeviceSettings)
	if err != nil {
		return err
	}

	// Unmarshal CosDevicePolicy
	var cosDevicePolicy pmpb.CosDevicePolicy
	err = proto.Unmarshal(cosDevicePolicyBytes, &cosDevicePolicy)
	if err != nil {
		return err
	}

	// Store the relevant protobufs in memory.
	p.autoUpdateSettings = chromeDeviceSettings.GetAutoUpdateSettings()
	p.metricsEnabled = chromeDeviceSettings.GetMetricsEnabled()
	p.cosDevicePolicy = &cosDevicePolicy

	return nil
}