Added get_status api in policy manager

BUG=b/222726350
TEST=presubmit
RELEASE_NOTE=Added get_status api in device_policy_manager

Change-Id: I5b4b3f4e27adce4b663bbe0fa74b6a07351a1a01
diff --git a/cmd/getstatus/getstatus.go b/cmd/getstatus/getstatus.go
new file mode 100644
index 0000000..b933f8e
--- /dev/null
+++ b/cmd/getstatus/getstatus.go
@@ -0,0 +1,38 @@
+// Copyright 2022 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 getstatus
+
+import (
+	"fmt"
+
+	"github.com/golang/glog"
+	"github.com/golang/protobuf/proto"
+	"policy-manager/pkg/imgstatus"
+)
+
+const (
+	lsbFile            = "/etc/lsb-release"
+	instanceConfigFile = "/var/lib/devicesettings/instance_config"
+)
+
+// HandleGetStatusCmd prints the status of the isntance to stdout.
+func HandleGetStatusCmd() {
+	status, err := imgstatus.GetStatus(lsbFile, instanceConfigFile)
+	if err != nil {
+		glog.Exitf("Error fetching instance status: %v", err)
+	}
+
+	fmt.Println(proto.MarshalTextString(status))
+}
diff --git a/main.go b/main.go
index 094b819..3c11bb9 100644
--- a/main.go
+++ b/main.go
@@ -39,6 +39,7 @@
 	"flag"
 	"fmt"
 
+	"policy-manager/cmd/getstatus"
 	"policy-manager/cmd/monitor"
 	"policy-manager/cmd/showconfig"
 
@@ -51,10 +52,14 @@
 // Command line flags for show_config mode.
 var showConfigModeFlags *flag.FlagSet
 
+// Command line flags for get_status mode.
+var getStatusModeFlags *flag.FlagSet
+
 const (
 	// Command names.
 	monitorModeCommand = "monitor"
 	showConfigCommand  = "show_config"
+	getStatusCommand   = "get_status"
 )
 
 func init() {
@@ -70,6 +75,12 @@
 		fmt.Println("Prints the current device policy for this instance to stdout.")
 	}
 
+	// Flags for get_status.
+	getStatusModeFlags = flag.NewFlagSet(getStatusCommand, flag.ExitOnError)
+	getStatusModeFlags.Usage = func() {
+		fmt.Println("Prints the status of updates available to stdout.")
+	}
+
 	// Log to stderr only by default
 	flag.Set("logtostderr", "true")
 }
@@ -83,8 +94,8 @@
 		flag.PrintDefaults()
 		fmt.Println("")
 
-		fmt.Printf("Supported commands:\n%s,%s\n",
-			monitorModeCommand, showConfigCommand)
+		fmt.Printf("Supported commands:\n%s,%s,%s\n",
+			monitorModeCommand, showConfigCommand, getStatusCommand)
 	}
 
 	// Parse the glog flags.
@@ -105,6 +116,8 @@
 		monitor.HandleMonitorCmd()
 	case showConfigCommand:
 		showconfig.HandleShowConfigCmd()
+	case getStatusCommand:
+		getstatus.HandleGetStatusCmd()
 	default:
 		glog.Errorf("Unrecognized command: %s\n", command)
 		flag.Usage()
diff --git a/pkg/imgstatus/imgstatus.go b/pkg/imgstatus/imgstatus.go
new file mode 100644
index 0000000..01d5c7e
--- /dev/null
+++ b/pkg/imgstatus/imgstatus.go
@@ -0,0 +1,218 @@
+// Copyright 2022 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.
+
+// Implementation of imgstatus which returns the update status of the instance
+
+package imgstatus
+
+import (
+	"bytes"
+	"fmt"
+	"github.com/golang/protobuf/proto"
+	"io"
+	"io/ioutil"
+	"policy-manager/pkg/configfetcher"
+	"policy-manager/pkg/devicepolicy"
+	"policy-manager/pkg/updateengine"
+	"policy-manager/protos"
+	"regexp"
+	"strconv"
+	"strings"
+)
+
+// Constants related to lsb-release file.
+const (
+	lsbChromeOSReleaseKey          = "CHROMEOS_RELEASE_VERSION"
+	lsbChromeOSReleaseChannelKey   = "CHROMEOS_RELEASE_TRACK"
+	lsbChromeOSReleaseMilestoneKey = "CHROMEOS_RELEASE_CHROME_MILESTONE"
+	updateEngineCmd                = "update_engine_client"
+)
+
+// parseKeyValuePairs parses text with the form
+//      key1=value1
+//      key2=value2
+//      ...
+// into a map.
+// All keys must have the format: "[a-zA-Z_]*".
+// Any non-conforming lines are ignored.
+func parseKeyValuePairs(reader io.Reader) (map[string]string, error) {
+	content, err := ioutil.ReadAll(reader)
+	if err != nil {
+		return nil, err
+	}
+
+	pairs := make(map[string]string)
+
+	// Regexp for matching key-value pairs.
+	matchPair := regexp.MustCompile(`[a-zA-Z_]*=.*$`)
+
+	// Extract release version string.
+	lines := strings.Split(string(content), "\n")
+	for _, line := range lines {
+		line = strings.TrimSpace(line)
+		if matchPair.MatchString(line) {
+			pair := strings.Split(line, "=")
+			key := strings.TrimSpace(pair[0])
+			val := strings.TrimSpace(pair[1])
+			pairs[key] = val
+		}
+	}
+
+	return pairs, nil
+}
+
+// convertOSReleaseChannel takes in a release channel string from the
+// lsb-release file and converts it into a ReleaseChannel type.
+// If the release channel string cannot be parsed,
+// RELEASE_CHANNEL_UNSPECIFIED will be returned.
+func convertOSReleaseChannel(channelString string) protos.ReleaseChannel {
+	// Regexp for matching release channel strings.
+	// Release channel strings have the format "<name>-channel". For
+	// example, "dev-channel".
+	pattern := `^([a-z]{1,}-channel)$`
+	matchVersionString := regexp.MustCompile(pattern)
+
+	// Check we have a valid release channel string.
+	if matchVersionString.MatchString(channelString) {
+		ch := strings.ToUpper(strings.Split(channelString, "-")[0])
+		if c, ok := protos.ReleaseChannel_value[ch]; ok {
+			return protos.ReleaseChannel(c)
+		}
+	}
+
+	// Unable to parse release channel.
+	return protos.ReleaseChannel_RELEASE_CHANNEL_UNSPECIFIED
+}
+
+// parseOSInfoFromLSB returns the current OS version and release channel as
+// provided by the lsb-release file. An error is returned only if the lsb file
+// cannot be parsed.
+func parseOSInfoFromLSB(lsbFile string) (*protos.OSVersion, error) {
+	content, err := ioutil.ReadFile(lsbFile)
+	if err != nil {
+		return nil, err
+	}
+
+	// Parse lsb-release file as key value pairs.
+	pairs, err := parseKeyValuePairs(bytes.NewReader(content))
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse '%s': %s", lsbFile, err)
+	}
+
+	osVersion := new(protos.OSVersion)
+
+	// Get release version string.
+	if version, ok := pairs[lsbChromeOSReleaseKey]; ok {
+		osVersion.VersionString = proto.String(version)
+	}
+
+	// Get the release channel.
+	if channel, ok := pairs[lsbChromeOSReleaseChannelKey]; ok {
+		ch := convertOSReleaseChannel(channel)
+		osVersion.Channel = &ch
+	}
+
+	// Get release milestone.
+	if milestone, ok := pairs[lsbChromeOSReleaseMilestoneKey]; ok {
+		if n, err := strconv.ParseUint(milestone, 10, 32); err == nil {
+			osVersion.Milestone = proto.Uint32(uint32(n))
+		}
+	}
+	return osVersion, nil
+}
+
+// getStatusFromUpdateEngine returns the protos.InstanceStatus protobuf
+// with the update status filled in.
+func getStatusFromUpdateEngine(instanceConfigFile string, updateEngineCmd string) (*protos.InstanceStatus, error) {
+	status := new(protos.InstanceStatus)
+
+	// Check if update engine is disabled.
+	config, err := devicepolicy.GetInstanceConfig(instanceConfigFile)
+	if err != nil {
+		return nil, fmt.Errorf("error fetching instance config: %v", err)
+	}
+
+	// Do not check update engine status if it is disabled.
+	if *config.UpdateStrategy == "update_disabled" {
+		status.UpdateStatus = proto.String("UPDATE_ENGINE_DISABLED")
+		return status, nil
+	}
+
+	// Get response from update engine.
+	response, err := updateengine.UpdateEngineStatus(updateEngineCmd)
+	if err != nil {
+		return nil, err
+	}
+
+	pairs, err := parseKeyValuePairs(bytes.NewReader(response))
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse: %s", err)
+	}
+
+	// Set update status.
+	if pairs["CURRENT_OP"] == "" {
+		return nil, fmt.Errorf("failed to fetch update status. It is possible that update_engine has not fetched the status yet.")
+	}
+	status.UpdateStatus = proto.String(pairs["CURRENT_OP"])
+
+	// Set new OS version.
+	status.NewOsVersion = new(protos.OSVersion)
+	status.NewOsVersion.VersionString = proto.String(pairs["NEW_VERSION"])
+
+	// Set last update check timestamp.
+	lastTime, err := strconv.ParseInt(pairs["LAST_CHECKED_TIME"], 10, 64)
+	if err != nil {
+		return nil, fmt.Errorf("unable to parse LAST_CHECKED_TIME: %s", err)
+	}
+	status.UpdateCheckTimestamp = proto.Int64(lastTime)
+
+	return status, nil
+}
+
+// GetStatus() returns the current status of the instance as given by
+// /etc/lsb-release file and update-engine.
+// If an error occured, the error will be returned along with a status protobuf
+// where the error field is set to the message describing the error. If the
+// error field is set, all other fields other than the InstanceId are invalid.
+// Note that this function must always return a status with InstanceId field
+// populated regardless of whether an error occurred or not.
+func GetStatus(lsbFile string, instanceConfigFile string) (*protos.InstanceStatus, error) {
+	status := new(protos.InstanceStatus)
+	var instanceID uint64
+	id, err := configfetcher.GetInstanceID()
+	if err == nil {
+		instanceID = id
+	}
+	status.InstanceId = proto.Uint64(instanceID)
+
+	// Get OS version from lsb-release file.
+	osVersion, err := parseOSInfoFromLSB(lsbFile)
+	if err != nil {
+		status.Error = proto.String(err.Error())
+		return status, err
+	}
+
+	status.OsVersion = osVersion
+
+	// Get update status from update-engine via update_engine_client.
+	updateStatus, err := getStatusFromUpdateEngine(instanceConfigFile, updateEngineCmd)
+	if err != nil {
+		status.UpdateStatus = nil
+		status.Error = proto.String(err.Error())
+		return status, err
+	}
+
+	proto.Merge(status, updateStatus)
+	return status, nil
+}
diff --git a/pkg/imgstatus/imgstatus_test.go b/pkg/imgstatus/imgstatus_test.go
new file mode 100644
index 0000000..fe9bd70
--- /dev/null
+++ b/pkg/imgstatus/imgstatus_test.go
@@ -0,0 +1,239 @@
+// Copyright 2022 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.
+
+// Unit tests for the imgstatus implementation.
+package imgstatus
+
+import (
+	"os"
+	"policy-manager/pkg/sysapi"
+	"policy-manager/protos"
+	"reflect"
+	"strings"
+	"testing"
+
+	"github.com/golang/protobuf/proto"
+)
+
+const (
+	// Test LSB file.
+	testLSB = `CHROMEOS_AUSERVER=https://tools.google.com/service/update2
+                CHROMEOS_BOARD_APPID={76E245CF-C0D0-444D-BA50-36739C18EB00}
+                CHROMEOS_DEVSERVER=
+                CHROMEOS_RELEASE_APPID={90F229CE-83E2-4FAF-8479-E368A34938B1}
+                CHROMEOS_RELEASE_BOARD=lakitu-signed-prempkeys
+                CHROMEOS_RELEASE_BRANCH_NUMBER=0
+                CHROMEOS_RELEASE_BUILD_NUMBER=16511
+                CHROMEOS_RELEASE_BUILD_TYPE=Official Build
+                CHROMEOS_RELEASE_CHROME_MILESTONE=89
+                CHROMEOS_RELEASE_DESCRIPTION=16511.0.0 (Official Build) dev-channel lakitu
+                CHROMEOS_RELEASE_NAME=Chrome OS
+                CHROMEOS_RELEASE_PATCH_NUMBER=0
+                CHROMEOS_RELEASE_TRACK=dev-channel
+                CHROMEOS_RELEASE_VERSION=16511.0.0
+                GOOGLE_RELEASE=16511.0.0
+                HWID_OVERRIDE=LAKITU`
+
+	testDevBuildLSB = `
+                CHROMEOS_RELEASE_BOARD=lakitu
+                CHROMEOS_DEVSERVER=http://foo.google.com:8080
+                GOOGLE_RELEASE=7520.3.2015_10_12_1443
+                CHROMEOS_RELEASE_BUILD_NUMBER=7520
+                CHROMEOS_RELEASE_BRANCH_NUMBER=3
+                CHROMEOS_RELEASE_CHROME_MILESTONE=89
+                CHROMEOS_RELEASE_PATCH_NUMBER=2015_10_12_1443
+                CHROMEOS_RELEASE_TRACK=developer-build
+                CHROMEOS_RELEASE_DESCRIPTION=7520.3.2015_10_12_1443 (Developer Build)
+                CHROMEOS_RELEASE_NAME=Chromium OS
+                CHROMEOS_RELEASE_BUILD_TYPE=Developer Build
+                CHROMEOS_RELEASE_VERSION=7520.3.2015_10_12_1443
+                CHROMEOS_AUSERVER=http://foo.google.com:8080/update`
+
+	// Dummy instance ID.
+	testInstanceID = uint64(12031993)
+	tmpLsbFile     = "/tmp/lsb-release"
+	tmpLsbFilePerm = os.FileMode(0644)
+)
+
+// TestParseKeyValuePairs tests that we can correctly parse /etc/lsb-release
+// into a map.
+func TestParseKeyValuePairs(t *testing.T) {
+	tests := []struct {
+		name  string
+		want  map[string]string
+		input string
+	}{
+		{
+			"ParseLSBFile",
+			map[string]string{
+				"CHROMEOS_AUSERVER":                 "https://tools.google.com/service/update2",
+				"CHROMEOS_BOARD_APPID":              "{76E245CF-C0D0-444D-BA50-36739C18EB00}",
+				"CHROMEOS_DEVSERVER":                "",
+				"CHROMEOS_RELEASE_APPID":            "{90F229CE-83E2-4FAF-8479-E368A34938B1}",
+				"CHROMEOS_RELEASE_BOARD":            "lakitu-signed-prempkeys",
+				"CHROMEOS_RELEASE_BRANCH_NUMBER":    "0",
+				"CHROMEOS_RELEASE_BUILD_NUMBER":     "16511",
+				"CHROMEOS_RELEASE_BUILD_TYPE":       "Official Build",
+				"CHROMEOS_RELEASE_CHROME_MILESTONE": "89",
+				"CHROMEOS_RELEASE_DESCRIPTION":      "16511.0.0 (Official Build) dev-channel lakitu",
+				"CHROMEOS_RELEASE_NAME":             "Chrome OS",
+				"CHROMEOS_RELEASE_PATCH_NUMBER":     "0",
+				"CHROMEOS_RELEASE_TRACK":            "dev-channel",
+				"CHROMEOS_RELEASE_VERSION":          "16511.0.0",
+				"GOOGLE_RELEASE":                    "16511.0.0",
+				"HWID_OVERRIDE":                     "LAKITU",
+			},
+			testLSB,
+		},
+		{
+			"PoorlyFormattedPairs",
+			map[string]string{
+				"apple_A":    "a",
+				"BAnaNA_b_":  "b",
+				"cantaloupe": "",
+			},
+			`
+                                apple_A= a
+                                BAnaNA_b_ =b
+                                    cantaloupe=
+                        `,
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			mockReader := strings.NewReader(test.input)
+			pairs, err := parseKeyValuePairs(mockReader)
+			if err != nil {
+				t.Errorf("got unexpected error %v", err)
+			} else if !reflect.DeepEqual(pairs, test.want) {
+				t.Errorf("got %v, want %v", pairs, test.want)
+			}
+		})
+	}
+}
+
+// TestConvertOSReleaseChannel tests that we can correctly convert the string
+// representation of release channel into the protobuf definition.
+func TestConvertOSReleaseChannel(t *testing.T) {
+	tests := []struct {
+		name  string
+		want  protos.ReleaseChannel
+		input string
+	}{
+		{
+			"DevChannelTest",
+			protos.ReleaseChannel_DEV,
+			"dev-channel",
+		},
+		{
+			"BetaChannelTest",
+			protos.ReleaseChannel_BETA,
+			"beta-channel",
+		},
+		{
+			"StableChannelTest",
+			protos.ReleaseChannel_STABLE,
+			"stable-channel",
+		},
+		{
+			"DeveloperBuildTest",
+			protos.ReleaseChannel_RELEASE_CHANNEL_UNSPECIFIED,
+			"developer-build",
+		},
+		{
+			"BadChannelTest",
+			protos.ReleaseChannel_RELEASE_CHANNEL_UNSPECIFIED,
+			"abc-chan",
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			result := convertOSReleaseChannel(test.input)
+			if !reflect.DeepEqual(test.want, result) {
+				t.Errorf("got %s, want %s",
+					result.String(),
+					test.want.String())
+			}
+		})
+	}
+}
+
+// TestParseOSInfoFromLSB tests that we can correctly get the OS version and
+// release channel from a lsb-release file.
+func TestParseOSInfoFromLSB(t *testing.T) {
+	devChannel := protos.ReleaseChannel_DEV
+	unknownChannel := protos.ReleaseChannel_RELEASE_CHANNEL_UNSPECIFIED
+
+	tests := []struct {
+		name       string
+		want       *protos.OSVersion
+		lsbContent string
+		expectErr  bool
+	}{
+		{
+			"GoodLSBFile",
+			&protos.OSVersion{
+				VersionString: proto.String("16511.0.0"),
+				Milestone:     proto.Uint32(uint32(89)),
+				Channel:       &devChannel,
+			},
+			testLSB,
+			false,
+		},
+		{
+			"OnlyReleaseVersion",
+			&protos.OSVersion{
+				VersionString: proto.String("16511.0.0"),
+			},
+			"CHROMEOS_RELEASE_VERSION=16511.0.0",
+			false,
+		},
+		{
+			"BadLSBFile",
+			&protos.OSVersion{},
+			"",
+			false,
+		},
+		{
+			"DevBuildLSBFile",
+			&protos.OSVersion{
+				VersionString: proto.String("7520.3.2015_10_12_1443"),
+				Milestone:     proto.Uint32(uint32(89)),
+				Channel:       &unknownChannel,
+			},
+			testDevBuildLSB,
+			false,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			sysapi.AtomicWriteFile(tmpLsbFile, []byte(test.lsbContent), tmpLsbFilePerm)
+			defer os.RemoveAll(tmpLsbFile)
+
+			result, err := parseOSInfoFromLSB(tmpLsbFile)
+
+			if err == nil && test.expectErr {
+				t.Errorf("test passed, want error")
+			} else if err == nil && !proto.Equal(test.want, result) {
+				t.Errorf("got %s, want %s",
+					result.String(),
+					test.want.String())
+			} else if err != nil && !test.expectErr {
+				t.Errorf("got unexpected error %v", err)
+			}
+		})
+	}
+}
diff --git a/pkg/policyenforcer/policy_enforcer.go b/pkg/policyenforcer/policy_enforcer.go
index 51c03e1..4cfbc15 100644
--- a/pkg/policyenforcer/policy_enforcer.go
+++ b/pkg/policyenforcer/policy_enforcer.go
@@ -115,11 +115,11 @@
 }
 
 // GetServiceStatus checks whether the services are running and returns their
-// status in a InstanceStatus proto. Failure to check one service will result
+// status in a ServiceStatus proto. Failure to check one service will result
 // in a missing field in the returned proto, and will not affect checking the
 // other service's status.
-func (client *PolicyEnforcer) GetServiceStatus(serviceMonitor map[string]string) (*protos.InstanceStatus, error) {
-	instanceStatus := new(protos.InstanceStatus)
+func (client *PolicyEnforcer) GetServiceStatus(serviceMonitor map[string]string) (*protos.ServiceStatus, error) {
+	ServiceStatus := new(protos.ServiceStatus)
 
 	statusErr := false
 	isRunning, err := client.systemdClient.IsUnitActiveRunning(serviceMonitor["metricsService"])
@@ -127,7 +127,7 @@
 		glog.Error(err)
 		statusErr = true
 	} else {
-		instanceStatus.Metrics = proto.Bool(isRunning)
+		ServiceStatus.Metrics = proto.Bool(isRunning)
 	}
 
 	isRunning, err = client.systemdClient.IsUnitActiveRunning(serviceMonitor["updateService"])
@@ -137,9 +137,9 @@
 
 	} else {
 		if isRunning {
-			instanceStatus.UpdateEngine = proto.String("")
+			ServiceStatus.UpdateEngine = proto.String("")
 		} else {
-			instanceStatus.UpdateEngine = proto.String(updateDisabledStrategy)
+			ServiceStatus.UpdateEngine = proto.String(updateDisabledStrategy)
 		}
 	}
 
@@ -148,7 +148,7 @@
 		glog.Error(err)
 		statusErr = true
 	} else {
-		instanceStatus.Logging = proto.Bool(isRunning)
+		ServiceStatus.Logging = proto.Bool(isRunning)
 	}
 
 	isRunning, err = client.systemdClient.IsUnitActiveRunning(serviceMonitor["monitoringService"])
@@ -157,11 +157,11 @@
 		statusErr = true
 
 	} else {
-		instanceStatus.Monitoring = proto.Bool(isRunning)
+		ServiceStatus.Monitoring = proto.Bool(isRunning)
 	}
 
 	if statusErr {
-		return instanceStatus, errors.New("unable to get service status")
+		return ServiceStatus, errors.New("unable to get service status")
 	}
-	return instanceStatus, nil
+	return ServiceStatus, nil
 }
diff --git a/pkg/policyenforcer/policy_enforcer_test.go b/pkg/policyenforcer/policy_enforcer_test.go
index 15ca5f0..b5da283 100644
--- a/pkg/policyenforcer/policy_enforcer_test.go
+++ b/pkg/policyenforcer/policy_enforcer_test.go
@@ -266,7 +266,7 @@
 		checkLoggingErr          error
 		isMonitoring             bool
 		checkMonitoringErr       error
-		expectedStatus           *protos.InstanceStatus
+		expectedStatus           *protos.ServiceStatus
 		expectErr                bool
 	}{
 		{
@@ -279,7 +279,7 @@
 			checkLoggingErr:          nil,
 			isMonitoring:             false,
 			checkMonitoringErr:       nil,
-			expectedStatus: &protos.InstanceStatus{
+			expectedStatus: &protos.ServiceStatus{
 				UpdateEngine: proto.String("update_disabled"),
 				Metrics:      proto.Bool(false),
 				Logging:      proto.Bool(false),
@@ -297,7 +297,7 @@
 			checkLoggingErr:          nil,
 			isMonitoring:             true,
 			checkMonitoringErr:       nil,
-			expectedStatus: &protos.InstanceStatus{
+			expectedStatus: &protos.ServiceStatus{
 				UpdateEngine: proto.String(""),
 				Metrics:      proto.Bool(true),
 				Logging:      proto.Bool(true),
@@ -315,7 +315,7 @@
 			checkLoggingErr:          nil,
 			isMonitoring:             false,
 			checkMonitoringErr:       nil,
-			expectedStatus: &protos.InstanceStatus{
+			expectedStatus: &protos.ServiceStatus{
 				UpdateEngine: proto.String(""),
 				Logging:      proto.Bool(true),
 				Monitoring:   proto.Bool(false),
@@ -332,7 +332,7 @@
 			checkLoggingErr:          nil,
 			isMonitoring:             false,
 			checkMonitoringErr:       nil,
-			expectedStatus: &protos.InstanceStatus{
+			expectedStatus: &protos.ServiceStatus{
 				Metrics:    proto.Bool(true),
 				Logging:    proto.Bool(true),
 				Monitoring: proto.Bool(false),
@@ -349,7 +349,7 @@
 			checkLoggingErr:          errors.New("error"),
 			isMonitoring:             true,
 			checkMonitoringErr:       nil,
-			expectedStatus: &protos.InstanceStatus{
+			expectedStatus: &protos.ServiceStatus{
 				UpdateEngine: proto.String(""),
 				Metrics:      proto.Bool(true),
 				Monitoring:   proto.Bool(true),
@@ -366,7 +366,7 @@
 			checkLoggingErr:          nil,
 			isMonitoring:             false,
 			checkMonitoringErr:       errors.New("error"),
-			expectedStatus: &protos.InstanceStatus{
+			expectedStatus: &protos.ServiceStatus{
 				UpdateEngine: proto.String(""),
 				Metrics:      proto.Bool(true),
 				Logging:      proto.Bool(true),
diff --git a/pkg/updateengine/testdata/update_engine.sh b/pkg/updateengine/testdata/update_engine.sh
new file mode 100755
index 0000000..3363bc0
--- /dev/null
+++ b/pkg/updateengine/testdata/update_engine.sh
@@ -0,0 +1,30 @@
+#!/bin/bash
+
+# Copyright 2022 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.
+
+# This is a fake implementation of update-engine service
+
+main() {
+    # perform the following action when the command is
+    # update-engine --status
+    if [[ $1 == "--status" ]]; then
+        echo "CURRENT_OP=UPDATE_STATUS_UPDATED_NEED_REBOOT\
+        \nLAST_CHECKED_TIME=1646430409\nNEW_VERSION=0.0.0.0"
+    else
+        exit 1
+    fi
+}
+
+main $@
diff --git a/pkg/updateengine/update_engine_client.go b/pkg/updateengine/update_engine_client.go
new file mode 100644
index 0000000..9855d7b
--- /dev/null
+++ b/pkg/updateengine/update_engine_client.go
@@ -0,0 +1,41 @@
+// Copyright 2022 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 updateengine
+
+import (
+	"bytes"
+	"fmt"
+	"github.com/golang/glog"
+	"os/exec"
+)
+
+// UpdateEngineStatus is used to get the status of update engine if auto updates
+// are enabled in the instance.
+func UpdateEngineStatus(updateEngineCmd string) ([]byte, error) {
+	cmd := exec.Command(updateEngineCmd, "--status")
+	var stdout, stderr bytes.Buffer
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+
+	if err := cmd.Run(); err != nil {
+		return nil, fmt.Errorf("unable to get update-engine status: %s", err)
+	}
+
+	if stderr.String() != "" {
+		glog.Errorf("error from update_engine_client: %v", stderr.String())
+	}
+
+	return stdout.Bytes(), nil
+}
diff --git a/pkg/updateengine/update_engine_client_test.go b/pkg/updateengine/update_engine_client_test.go
new file mode 100644
index 0000000..00fc1e4
--- /dev/null
+++ b/pkg/updateengine/update_engine_client_test.go
@@ -0,0 +1,34 @@
+// Copyright 2022 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.
+
+// Unit tests for the TestUpdateEngineStatus implementation.
+package updateengine
+
+import (
+	"reflect"
+	"testing"
+)
+
+const tmpUpdateEngineCmd = "testdata/update_engine.sh"
+
+func TestUpdateEngineStatus(t *testing.T) {
+	want := []byte("CURRENT_OP=UPDATE_STATUS_UPDATED_NEED_REBOOT\nLAST_CHECKED_TIME=1646430409\nNEW_VERSION=0.0.0.0")
+	got, err := UpdateEngineStatus(tmpUpdateEngineCmd)
+	if err != nil {
+		t.Errorf("got unexpected error: %v", err)
+	}
+	if reflect.DeepEqual(got, want) {
+		t.Errorf("got: %s, want: %s", got, want)
+	}
+}
diff --git a/protos/instance_status.proto b/protos/instance_status.proto
index 72ca8be..f268b2c 100644
--- a/protos/instance_status.proto
+++ b/protos/instance_status.proto
@@ -1,4 +1,4 @@
-// Copyright 2021 Google LLC
+// Copyright 2022 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -16,7 +16,40 @@
 syntax = "proto2";
 
 package protos;
-// option go_package = "policy_manager/protos";
+
+// ReleaseChannel specifies the Chrome OS release channel that the image is on.
+enum ReleaseChannel {
+  // Unknown release channel.
+  RELEASE_CHANNEL_UNSPECIFIED = 0;
+
+  // Stable channel.
+  STABLE = 1;
+
+  // Beta channel.
+  BETA = 2;
+
+  // Dev channel.
+  DEV = 3;
+
+  reserved 4;
+}
+
+// OSVersion contains all the fields identifying a particular version of the
+// Chrome OS.
+message OSVersion {
+  // Release version string. For example: "16511.0.0"
+  // Chrome OS uses the following version format:
+  // <TIP_BUILD>.<BRANCH_BUILD>.<BRANCH_BRANCH_BUILD>
+  optional string version_string = 1;
+
+  // Milestone number for Chrome. For example: 89
+  optional uint32 milestone = 2;
+
+  // Release channel of the OS.
+  optional ReleaseChannel channel = 3;
+
+  // Next number: 4
+}
 
 // InstanceStatus contains all the fields related to the current status of the
 // instance.
@@ -24,26 +57,25 @@
   // GCE instance ID.
   optional uint64 instance_id = 1;
 
-  // update_engine is set to true if the update engine systemd
-  // service is active.
-  optional string update_engine = 2;
+  // Current version of the image.
+  optional OSVersion os_version = 2;
 
-  // metrics is set to true if the crash reporter systemd
-  // service is active.
-  optional bool metrics = 3;
+  // The current update status as reported by the update_engine.
+  optional string update_status = 3;
 
-  // logging is set to true if the logging systemd service
-  // is active and running.
-  optional bool logging = 4;
+  // New version delivered by the update.
+  optional OSVersion new_os_version = 4;
 
-  // monitoring is set to true if the monitoring systemd service
-  // is active and running.
-  optional bool monitoring = 5;
+  // Unix timestamp of last reboot.
+  optional int64 reboot_timestamp = 5;
+
+  // Unix timestamp fo last update check.
+  optional int64 update_check_timestamp = 6;
 
   // If an error occurred while collecting the instance status, this field will
   // be set with the error message and all other fields except the instance_id
   // are undefined.
-  optional string error = 6;
+  optional string error = 7;
 
-  // Next number: 7
+  // Next number: 8
 }
diff --git a/protos/service_status.proto b/protos/service_status.proto
new file mode 100644
index 0000000..8147ddd
--- /dev/null
+++ b/protos/service_status.proto
@@ -0,0 +1,49 @@
+// 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.
+
+// Protobuf definitions for reporting instance status.
+syntax = "proto2";
+
+package protos;
+// option go_package = "policy_manager/protos";
+
+// ServiceStatus contains all the fields related to the current status of the
+// instance.
+message ServiceStatus {
+  // GCE instance ID.
+  optional uint64 instance_id = 1;
+
+  // update_engine is set to true if the update engine systemd
+  // service is active.
+  optional string update_engine = 2;
+
+  // metrics is set to true if the crash reporter systemd
+  // service is active.
+  optional bool metrics = 3;
+
+  // logging is set to true if the logging systemd service
+  // is active and running.
+  optional bool logging = 4;
+
+  // monitoring is set to true if the monitoring systemd service
+  // is active and running.
+  optional bool monitoring = 5;
+
+  // If an error occurred while collecting the instance status, this field will
+  // be set with the error message and all other fields except the instance_id
+  // are undefined.
+  optional string error = 6;
+
+  // Next number: 7
+}