// 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"
	"io/ioutil"
	"os"
	"time"

	"policy-manager/pkg/sysapi"
	pmpb "policy-manager/protos"
	pb "policy-manager/third_party/chromium/devicepolicyproto"

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

// 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 {
	metricsEnabled *pb.MetricsEnabledProto
	// InstanceConfig is an InstanceConfigProto that manages the device policy file.
	instanceConfig *pmpb.InstanceConfig
}

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

// 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) (*pb.PolicyFetchResponse, error) {
	// Generate ChromeDeviceSettingsProto.
	chromeDeviceSettings := pb.ChromeDeviceSettingsProto{
		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(time.Now().Unix()),
		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

	instanceConfig := &pmpb.InstanceConfig{
		UpdateStrategy:      proto.String(config.GetUpdateStrategy()),
		MetricsEnabled:      proto.Bool(config.GetMetricsEnabled()),
		HealthMonitorConfig: proto.Clone(config.HealthMonitorConfig).(*pmpb.HealthMonitorConfig),
	}

	metricsEnabled := &pb.MetricsEnabledProto{
		MetricsEnabled: proto.Bool(config.GetMetricsEnabled()),
	}

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

	p.instanceConfig = instanceConfig
	p.metricsEnabled = metricsEnabled
	return policyChanged
}

// ToInstanceConfig converts the DevicePolicy into InstanceConfig proto
func (p *DevicePolicy) ToInstanceConfig(config *pmpb.InstanceConfig) {
	if config == nil {
		return
	}

	if p.instanceConfig != nil {
		config.UpdateStrategy = proto.String(p.instanceConfig.GetUpdateStrategy())
		config.MetricsEnabled = proto.Bool(p.instanceConfig.GetMetricsEnabled())
		config.HealthMonitorConfig = proto.Clone(p.instanceConfig.HealthMonitorConfig).(*pmpb.HealthMonitorConfig)
	}
}

// WriteFile writes the policy to disk. The policy will atomically replace the existing policy.
func (p *DevicePolicy) SignAndWrite(key *rsa.PrivateKey, devicePolicyFile string, instanceConfigFile string) error {
	// Generate policy protobuf.
	policyFetchResponse, err := p.generateDevicePolicy(key)
	if err != nil {
		return nil
	}
	policyBytes, err := proto.Marshal(policyFetchResponse)
	if err != nil {
		return err
	}
	// Generate policy protobuf.
	instanceConfigBytes, err := proto.Marshal(p.instanceConfig)
	if err != nil {
		return err
	}
	// Write the policy file atomically.
	if err := sysapi.AtomicWriteFile(devicePolicyFile, policyBytes, configFilePerm); err != nil {
		return err
	}
	// Write the policy file atomically.
	if err := sysapi.AtomicWriteFile(instanceConfigFile, instanceConfigBytes, configFilePerm); err != nil {
		return err
	}
	return nil
}

// ReadPolicyFromDisk reads the device policy from disk.
// No error is returned for missing policy file.
func (p *DevicePolicy) ReadPolicyFromDisk(instanceConfigFile string) error {
	// Check if cos device policy file exists
	if _, err := os.Stat(instanceConfigFile); os.IsNotExist(err) {
		return nil
	}

	// Read the policy file and convert it to protobuf format.
	instanceConfigBytes, err := ioutil.ReadFile(instanceConfigFile)
	if err != nil {
		return err
	}

	// Unmarshal InstanceConfig
	var instanceConfig pmpb.InstanceConfig
	err = proto.Unmarshal(instanceConfigBytes, &instanceConfig)
	if err != nil {
		return err
	}

	// Store the relevant protobufs in memory.
	p.instanceConfig = &instanceConfig

	return nil
}