From f41e273340bc606127eeb7627fb918d7d8cc719c Mon Sep 17 00:00:00 2001
From: Dheyay <dheyay.desai@canonical.com>
Date: Tue, 5 Nov 2024 08:27:37 -0800
Subject: [PATCH 1/5] Added golang hook

---
 apt-hook/apthook/counts.go     | 223 +++++++++++++++++++++++++++++++++
 apt-hook/apthook/esm_counts.go |  62 +++++++++
 apt-hook/apthook/main.go       |  33 +++++
 apt-hook/apthook/rpc.go        |  98 +++++++++++++++
 apt-hook/apthook/types.go      |  48 +++++++
 apt-hook/json-hook.cc          |   3 +-
 6 files changed, 466 insertions(+), 1 deletion(-)
 create mode 100644 apt-hook/apthook/counts.go
 create mode 100644 apt-hook/apthook/esm_counts.go
 create mode 100644 apt-hook/apthook/main.go
 create mode 100644 apt-hook/apthook/rpc.go
 create mode 100644 apt-hook/apthook/types.go

diff --git a/apt-hook/apthook/counts.go b/apt-hook/apthook/counts.go
new file mode 100644
index 0000000000..7f67456d2e
--- /dev/null
+++ b/apt-hook/apthook/counts.go
@@ -0,0 +1,223 @@
+package apthook
+
+import "fmt"
+
+const (
+	cloudIDfile = "/run/cloud-init/cloud-id"
+	osReleaseFile = "/etc/os-release"
+)
+
+type SecurityCounts struct {
+	Standard int
+	ESMInfra int
+	ESMApps int
+}
+
+// Gets the ubuntu distro for the package
+func getDistroFromPackage(rpc *jsonRPC) string {
+	for _, pkg := range rpc.Params.Packages {
+		for _, origin := range pkg.Versions.Candidate.Origins {
+			if origin.Codename != "" {
+				return origin.Codename
+			}
+		}
+	}
+	return ""
+}
+
+// Verify the package is from specific origin
+func verifyOrigin(version jsonRPCPackageVersion, origin string) bool {
+	for _, packageOrigin := range version.Origins {
+		if packageOrigin.Origin == origin {
+			return true
+		}
+	}
+	return false
+}
+
+// Verify the package is from specific origin and archive
+func verifyOriginAndArchive(version jsonRPCPackageVersion, origin string, archive string) bool {
+	for _, packageOrigin := range version.Origins {
+		if packageOrigin.Origin == origin && packageOrigin.Archive == archive {
+			return true
+		}
+	}
+	return false
+}
+
+// Count security packages from apt stats
+func CountSecurityUpdates(rpc *jsonRPC) *SecurityCounts {
+	counts := &SecurityCounts{}
+	for _, pkg := range rpc.Params.Packages {
+		if pkg.Mode == "upgrade" {
+			// Check esm
+			if verifyOriginAndArchive(pkg.Versions.Install, "UbuntuESMApps", "-apps-security") {
+				counts.ESMApps++
+			}
+			// Checm esm-infra
+			if verifyOriginAndArchive(pkg.Versions.Install, "UbuntuESM", "-infra-security") {
+				counts.ESMInfra++
+			}
+			// Check standard
+			if verifyOriginAndArchive(pkg.Versions.Install, "UbuntuStandard", "-security") {
+				counts.Standard++
+			}
+		}
+	}
+	return counts
+}
+
+// CreateCountMessage generates a formatted message describing security updates
+func GetCountMessage(counts *SecurityCounts) string {
+	if counts.Standard == 0 && counts.ESMInfra == 0 && counts.ESMApps == 0 {
+		return ""
+	}
+
+	if counts.ESMInfra == 0 && counts.ESMApps == 0 {
+		if counts.Standard == 1 {
+			return fmt.Sprint("1 standard LTS security update")
+		}
+		if counts.Standard > 1 {
+			return fmt.Sprintf("%d standard LTS security updates", counts.Standard)
+		}
+	}
+
+	message := make([]string, 0, 3)
+
+	// Add standard count
+	if counts.Standard > 0 {
+		if counts.Standard == 1 {
+			message = append(message, "1 standard LTS security update")
+		} else {
+			message = append(message, fmt.Sprintf("%d standard LTS security updates", counts.Standard))
+		}
+	}
+
+	// Add ESM Infra count
+	if counts.ESMInfra > 0 {
+		if counts.ESMInfra == 1 {
+			message = append(message, "1 esm-infra security update")
+		} else {
+			message = append(message, fmt.Sprintf("%d esm-infra security updates", counts.ESMInfra))
+		}
+	}
+
+	// Add ESM Apps count
+	if counts.ESMApps > 0 {
+		if counts.ESMApps == 1 {
+			message = append(message, "1 esm-apps security update")
+		} else {
+			message = append(message, fmt.Sprintf("%d esm-apps security updates", counts.ESMApps))
+		}
+	}
+
+	switch len(message) {
+		case 1:
+			return message[0]
+		case 2:
+			return fmt.Sprintf("%s and %s", message[0], message[1])
+		case 3:
+			return fmt.Sprintf("%s, %s and %s", message[0], message[1], message[2])
+	}
+
+	return ""
+}
+
+func CollectProPackagesFromRPC(rpc *jsonRPC) []string {
+	var expiredPackages []string
+
+	proOrigins := []string{
+		"UbuntuESM",
+		"UbuntuESMApps",
+		"UbuntuCC",
+		"UbuntuCIS",
+		"UbuntuFIPS",
+		"UbuntuFIPSUpdates",
+		"UbuntuFIPSPreview",
+		"UbuntuRealtimeKernel",
+		"UbuntuROS",
+		"UbuntuROSUpdates",
+	}
+
+	for _, pkg := range rpc.Params.Packages {
+		if pkg.Mode == "upgrade" {
+			for _, origin := range proOrigins {
+				if verifyOrigin(pkg.Versions.Install, origin) {
+					expiredPackages = append(expiredPackages, pkg.Name)
+					break
+				}
+			}
+		}
+	}
+
+	return expiredPackages
+}
+
+// PrintExpiredProPackages logs messages about expired pro package status
+func PrintExpiredProPackages(packageNames []string) {
+	if len(packageNames) == 0 {
+		return
+	}
+
+	fmt.Println("The following packages will fail to download because your Ubuntu Pro subscription has expired:")
+
+	currLine := " "
+	for _, pkg := range packageNames {
+		if len(currLine)+1+len(pkg) >= 79 {
+			fmt.Println(currLine)
+			currLine = " "
+		}
+		currLine = currLine + " " + pkg
+	}
+	if len(currLine) > 1 {
+		fmt.Println(currLine)
+	}
+
+	fmt.Println("Renew your subscription or run `sudo pro detach` to remove these errors")
+}
+
+func getCloudID() string {
+	file, err := os.Open(cloudIDfile)
+	if err != nil {
+		return ""
+	}
+	defer file.Close()
+	contet, err = ioutil.ReadAll(file)
+	if err != nil {
+		return ""
+	}
+	fileContent = string(content)
+	if strings.Contains(fileContent, "aws") {
+		return "AWS"
+	} else if strings.Contains(fileContent, "gcp") {
+		return "GCP"
+	} else if strings.Contains(fileContent, "azure") {
+		return "Azure"
+	}
+	return ""
+}
+
+func getOSRelease() string {
+	file, err := os.Open(osReleaseFile)
+	if err != nil {
+		return ""
+	}
+	defer file.Close()
+
+	content, err = ioutil.ReadAll(file)
+	if err != nil {
+		return ""
+	}
+
+	fileContent = string(content)
+	if strings.Contains(fileContent, "xenial") {
+		return "XENIAL"
+	} else if strings.Contains(fileContent, "bionic") {
+		return "BIONIC"
+	}
+	return "NOT_ESM_INFRA"
+}
+
+func printLearnMoreContent(){
+
+}
diff --git a/apt-hook/apthook/esm_counts.go b/apt-hook/apthook/esm_counts.go
new file mode 100644
index 0000000000..364d93bc17
--- /dev/null
+++ b/apt-hook/apthook/esm_counts.go
@@ -0,0 +1,62 @@
+package apthook
+
+import (
+	"fmt"
+	"strings"
+)
+
+const (
+	systemPkgStatus = "/var/lib/dpkg/status"
+	esmPkgStatus	= "/var/lib/ubuntu-advantage/apt-esm/var/lib/dpkg/status"
+)
+
+type PackageStatus struct {
+	Name	string
+	Version string
+	Status  string
+	Source  string
+}
+
+func readPackages(filepath string) ([]PackageStatus, error) {
+	file, err := os.Open(filepath)
+	if err != nil {
+		return nil, error
+	}
+	defer file.Close()
+
+	var packages []PackageStatus
+	var curPackage PackageStatus
+
+	scanner := bufio.NewScanner(file)
+	parsingPackage := false
+
+	for scanner.Scan(){
+		line := strings.TrimSpace(scanner.Text())
+
+		if strings.HasPrefix(line, "Package:") {
+			if parsingPackage {
+				packages = append(packages, currentPackage)
+			}
+
+			curPackage = PackageStatus{}
+			curPackage.Name = strings.TrimSpace(strings.TrimPrefix(line, "Package:"))
+			parsingPackage = true
+		} else if strings.HasPrefix(line, "Version:") {
+			curPackage.Version = strings.TrimSpace(strings.TrimPrefix(line, "Version:"))
+		} else if strings.HasPrefix(line, "Status:") {
+			curPackage.Status = strings.TrimSpace(strings.TrimPrefix(line, "Status:"))
+		} else if strings.HasPrefix(line, "Source:") {
+			curPackage.Source = strings.TrimSpace(strings.TrimPrefix(line, "Source:"))
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return nil, err
+	}
+
+	if parsingPackage {
+		packages = append(packages, curPackage)
+	}
+
+	return packages, nil
+}
diff --git a/apt-hook/apthook/main.go b/apt-hook/apthook/main.go
new file mode 100644
index 0000000000..2b9c239966
--- /dev/null
+++ b/apt-hook/apthook/main.go
@@ -0,0 +1,33 @@
+package apthook
+
+import (
+	"fmt"
+	"strings"
+	"os"
+	"os/exec"
+)
+
+const (
+	methodStats = "org.debian.apt.hooks.install.statistics"
+	methodPrePrompt = "org.debian.apt.hooks.install.pre-prompt"
+	aptNewsFile = "/var/lib/ubuntu-advantage/messages/apt-news"
+	expiredNotice = "/var/lib/ubuntu-advantage/notices/5-contract_expired"
+)
+
+func main() {
+	// Get apt hook socker
+	// Make sure socket in not empty and exists
+	// jsonrpc handshake
+	// jsonrpc read rpc message
+	//
+	// if method status
+	// 	display security count message
+	// if method pre-prompt
+	//  get potential esm updates
+	//  call display functions from counts file
+	//
+	// Display apt news
+	// Display contract expiration notice
+	//
+	// jsonpc bye bye
+}
diff --git a/apt-hook/apthook/rpc.go b/apt-hook/apthook/rpc.go
new file mode 100644
index 0000000000..8efa87367a
--- /dev/null
+++ b/apt-hook/apthook/rpc.go
@@ -0,0 +1,98 @@
+package apthook
+
+importy (
+	"bufio"
+	"encoding/json"
+	"io"
+	"net"
+	"fmt"
+)
+
+const (
+	MethodHello      = "org.debian.apt.hooks.hello"
+	MethodBye		 = "org.debian.apt.hooks.bye"
+)
+
+type Connection struct {
+	conn net.Conn
+	reader *bufio.Reader
+}
+
+func (c *Connection) Close() error {
+	return c.conn.Close()
+}
+
+func (c *Connection) ReadMessage() (*jsonRPC, error) {
+	return readRPC(c.reader)
+}
+
+func (c *Connection) WriteResponse(version string, id int) error {
+	respo := fmt.Sprintf(`{"jsonrpc": "2.0", "id": %d, "result": {"version": "%s"}}`, id, version)
+	_, err := c.conn.Write([]byte(response + "\n\n"))
+	return err
+}
+
+func (c *Connection) Handshake() error {
+	msg, err := c.ReadMessage()
+	if err != nil {
+		return fmt.Errorf("reading handshake: %w", err)
+	}
+
+	if msg.Method != MethodHello {
+		return fmt.Errorf("expected hello method, got: %v", msg.Method)
+	}
+
+	if err := c.WriteResponse("0.2", 0); err != nil {
+		return fmt.Errorf("writing handshake response: %w", err)
+	}
+
+	return nil
+}
+
+func (c *Connection) Bye() error {
+	msg, err := c.ReadMessage()
+	if err != nil {
+		return fmt.Errorf("reading bye message: %w", err)
+	}
+
+	if msg.Method != MethodBye {
+		return fmt.Errorf("expected bye method, got: %v", msg.Method)
+	}
+
+	return nil
+}
+
+func NewConnection(f *os.File) (*Connection, error) {
+	conn, err := net.FileConn(f)
+	if err != nil {
+		return nil, err
+	}
+	return &Connection{
+		conn: conn,
+		reader: bufio.NewReader(conn),
+	}, nil
+}
+
+func readRPC(r *bufio.Reader) (*jsonRPC, error) {
+	line, err := r.ReadBytes('\n')
+	if err != nil && err != io.EOF {
+		return nil, err
+	}
+
+	var msg jsonRPC
+	if err := json.Unmarshal(line, &msg); err != nil {
+		return nul, fmt.Errorf("parsing json: %w", err)
+	}
+
+	emptyLine, _, err := r.ReadLine()
+
+	if err != nil {
+		return nul, fmt.Errorf("reading empty line: %w", err)
+	}
+
+	if string(emptyLine) != "" {
+		return nil, fmt.Errorf("unexpected line: %q (empty)", emptyLine)
+	}
+
+	return &msg, nil
+}
diff --git a/apt-hook/apthook/types.go b/apt-hook/apthook/types.go
new file mode 100644
index 0000000000..9b07694c3f
--- /dev/null
+++ b/apt-hook/apthook/types.go
@@ -0,0 +1,48 @@
+package apthook
+
+import (
+	"encoding/json"
+)
+
+type jsonRPC struct {
+	JsonRPC	string	`json:"jsonrpc"`
+	Method	string	`json:"method"`
+	Params	struct {
+		Command	        string   `json:"command"`
+		UnknownPackages	[]string `json:"unknown-packages"`
+		Packages		[]Package `json:"packages"`
+	} `json:"params"`
+}
+
+type Package struct {
+	ID           int    `json:"id"`
+	Name         string `json:"name"`
+	Architecture string `json:"architecture"`
+	Mode         string `json:"mode"`
+	Versions     Versions `json:"versions,omitempty"` // optional
+	Automatic    bool `json:"automatic,omitempty"`	// optional
+	Current	     string `json:"current,omitempty"`	// optional
+}
+
+type Versions struct {
+	Candidate	jsonRPCPackageVersion	`json:"candidate"`
+	Install		jsonRPCPPackageVersion	`json:"install"`
+	Current		jsonRPCPackageVersion	`json:"current"`
+}
+
+type jsonRPCPackageVersion struct {
+	ID           int    `json:"id"`
+	Version      string `json:"version"`
+	Architecture string `json:"architecture"`
+	Pin          int    `json:"pin"`
+	Origins      []Origin `json:"origins"`
+}
+
+type Origin struct {
+	Archive  string `json:"archive"`
+	Codename string `json:"codename"`
+	Version  string `json:"version"`
+	Origin   string `json:"origin"`
+	Label    string `json:"label"`
+	Site     string `json:"site"`
+}
diff --git a/apt-hook/json-hook.cc b/apt-hook/json-hook.cc
index 45650c5961..f9ba69dcca 100644
--- a/apt-hook/json-hook.cc
+++ b/apt-hook/json-hook.cc
@@ -10,6 +10,7 @@
 #include "json-hook.hh"
 #include "esm-counts.hh"
 
+// Written
 bool read_jsonrpc_request(std::istream &in, jsonrpc_request &req) {
     std::string msg_line;
     std::string empty_line;
@@ -193,7 +194,7 @@ bool collect_pro_packages_from_pre_prompt_json(json_object *pre_prompt, std::vec
             continue;
         }
         std::string package_mode(json_object_get_string(tmp));
-        
+
         has_key = json_object_object_get_ex(package, "name", &tmp);
         if (!has_key) {
             continue;

From 2224f90de2be116545d92fa55c01b83f634aea36 Mon Sep 17 00:00:00 2001
From: Dheyay <dheyay.desai@canonical.com>
Date: Tue, 12 Nov 2024 13:39:19 -0800
Subject: [PATCH 2/5] Add goland files

---
 apt-hook/apthook/counts.go     | 173 +++++++++++++++++++++++----------
 apt-hook/apthook/esm_counts.go |  12 +++
 apt-hook/apthook/types.go      |   8 ++
 3 files changed, 143 insertions(+), 50 deletions(-)

diff --git a/apt-hook/apthook/counts.go b/apt-hook/apthook/counts.go
index 7f67456d2e..5c6367f4e5 100644
--- a/apt-hook/apthook/counts.go
+++ b/apt-hook/apthook/counts.go
@@ -67,6 +67,75 @@ func CountSecurityUpdates(rpc *jsonRPC) *SecurityCounts {
 	return counts
 }
 
+func getCloudID() string {
+	file, err := os.Open(cloudIDfile)
+	if err != nil {
+		return ""
+	}
+	defer file.Close()
+	contet, err = ioutil.ReadAll(file)
+	if err != nil {
+		return ""
+	}
+	fileContent = string(content)
+	if strings.Contains(fileContent, "aws") {
+		return "AWS"
+	} else if strings.Contains(fileContent, "gcp") {
+		return "GCP"
+	} else if strings.Contains(fileContent, "azure") {
+		return "Azure"
+	}
+	return ""
+}
+
+// Process os release into a struct
+func parseOSRelease() OSRelease {
+	file, err := os.Open(osReleaseFile)
+	if err != nil {
+		return ""
+	}
+	defer file.Close()
+
+	raw_data := make(map[string]string)
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := scanner.text()
+		line_split := strings.SplitN(line, "=", 2)
+		if len(line_split) != 2 {
+			continue
+		}
+		key := strings.ToLower(line_split[0])
+		value := strings.Trim(line_split[1], `"`)
+
+		raw_data[key] = value
+	}
+
+	osRelease := OSRelease{
+		Name:             raw_data["NAME"],
+		VersionID:        raw_data["VERSION_ID"],
+		Version:          raw_data["VERSION"],
+		VersionCodename:  raw_data["VERSION_CODENAME"],
+		ID:               raw_data["ID"],
+	}
+
+	return osRelease
+}
+
+// Print package names from list
+func printPackageNames(packageNames []string) {
+	currLine := " "
+	for _, pkg := range packageNames {
+		if len(currLine)+1+len(pkg) >= 79 {
+			fmt.Println(currLine)
+			currLine = " "
+		}
+		currLine = currLine + " " + pkg
+	}
+	if len(currLine) > 1 {
+		fmt.Println(currLine)
+	}
+}
+
 // CreateCountMessage generates a formatted message describing security updates
 func GetCountMessage(counts *SecurityCounts) string {
 	if counts.Standard == 0 && counts.ESMInfra == 0 && counts.ESMApps == 0 {
@@ -123,6 +192,7 @@ func GetCountMessage(counts *SecurityCounts) string {
 	return ""
 }
 
+
 func CollectProPackagesFromRPC(rpc *jsonRPC) []string {
 	var expiredPackages []string
 
@@ -158,66 +228,69 @@ func PrintExpiredProPackages(packageNames []string) {
 	if len(packageNames) == 0 {
 		return
 	}
-
 	fmt.Println("The following packages will fail to download because your Ubuntu Pro subscription has expired:")
-
-	currLine := " "
-	for _, pkg := range packageNames {
-		if len(currLine)+1+len(pkg) >= 79 {
-			fmt.Println(currLine)
-			currLine = " "
-		}
-		currLine = currLine + " " + pkg
-	}
-	if len(currLine) > 1 {
-		fmt.Println(currLine)
-	}
-
+	printPackageNames(packageNames)
 	fmt.Println("Renew your subscription or run `sudo pro detach` to remove these errors")
 }
 
-func getCloudID() string {
-	file, err := os.Open(cloudIDfile)
-	if err != nil {
-		return ""
-	}
-	defer file.Close()
-	contet, err = ioutil.ReadAll(file)
-	if err != nil {
-		return ""
-	}
-	fileContent = string(content)
-	if strings.Contains(fileContent, "aws") {
-		return "AWS"
-	} else if strings.Contains(fileContent, "gcp") {
-		return "GCP"
-	} else if strings.Contains(fileContent, "azure") {
-		return "Azure"
+// Seperate function to check if os series is esm_infra
+func getESMInfraSeries() string {
+	osRelease := parseOSRelease()
+	if osRelease.VersionCodename == "xenial" {
+		return "XENIAL"
+	} else if osRelease.VersionCodename == "bionic" {
+		return "BIONIC"
 	}
-	return ""
+	return "NOT_ESM_INFRA"
 }
 
-func getOSRelease() string {
-	file, err := os.Open(osReleaseFile)
-	if err != nil {
-		return ""
-	}
-	defer file.Close()
+func printLearnMoreContent() {
+	cloudID := getCloudID()
+	esmInfraSeries := getOSRelease()
 
-	content, err = ioutil.ReadAll(file)
-	if err != nil {
-		return ""
+	if esmInfraSeries == "XENIAL" {
+		if cloudID == "Azure" {
+			fmt.Printf("Learn more about Ubuntu Pro for 16.04 on Azure at %s\n",
+				"https://ubuntu.com/16-04/azure")
+			return
+		} else {
+			fmt.Printf("Learn more about Ubuntu Pro for 16.04 at %s\n",
+				"https://ubuntu.com/16-04")
+			return
+		}
+	} else if esmInfraSeries == "BIONIC" {
+		if cloudID == "Azure" {
+			fmt.Printf("Learn more about Ubuntu Pro for 18.04 on Azure at %s\n",
+				"https://ubuntu.com/18-04/azure")
+			return
+		} else {
+			fmt.Printf("Learn more about Ubuntu Pro for 18.04 at %s\n",
+				"https://ubuntu.com/18-04")
+			return
+		}
+	} else {
+		if cloudID == "Azure" {
+			fmt.Printf("Learn more about Ubuntu Pro on Azure at %s\n",
+				"https://ubuntu.com/azure/pro")
+			return
+		} else if cloudID == "AWS" {
+			fmt.Printf("Learn more about Ubuntu Pro on AWS at %s\n",
+				"https://ubuntu.com/aws/pro")
+			return
+		} else if cloudID == "GCP" {
+			fmt.Printf("Learn more about Ubuntu Pro on GCP at %s\n",
+				"https://ubuntu.com/gcp/pro")
+			return
+		}
 	}
 
-	fileContent = string(content)
-	if strings.Contains(fileContent, "xenial") {
-		return "XENIAL"
-	} else if strings.Contains(fileContent, "bionic") {
-		return "BIONIC"
-	}
-	return "NOT_ESM_INFRA"
+	fmt.Printf("Learn more about Ubuntu Pro at %s\n",
+		"https://ubuntu.com/pro")
+	return
 }
 
-func printLearnMoreContent(){
-
+func printESMPackages(esmType ESMType, packageNames []string) {
+	fmt.Println("The following packages will fail to download because your Ubuntu Pro subscription has expired")
+	printPackageNames(packageNames)
+	fmt.Println("Renew your subscription or `sudo pro detach` to remove these errors")
 }
diff --git a/apt-hook/apthook/esm_counts.go b/apt-hook/apthook/esm_counts.go
index 364d93bc17..ec583f8315 100644
--- a/apt-hook/apthook/esm_counts.go
+++ b/apt-hook/apthook/esm_counts.go
@@ -9,6 +9,13 @@ const (
 	systemPkgStatus = "/var/lib/dpkg/status"
 	esmPkgStatus	= "/var/lib/ubuntu-advantage/apt-esm/var/lib/dpkg/status"
 )
+// Instead of esmPkgStatus
+// we need to use the following from var/lib/ubuntu-advantage/apt-esm/var/lib/apt/lists:
+// esmAppsSecurityList
+// esmAppsUpdatesList
+// esmInfraSecurityList
+// esmInfraUpdatesList
+
 
 type PackageStatus struct {
 	Name	string
@@ -17,6 +24,11 @@ type PackageStatus struct {
 	Source  string
 }
 
+// Func to get file pathes for esm packages
+// Based on os release codename
+
+// Packages need to be read and parsed from all above files
+// and then returned as a slice of PackageStatus
 func readPackages(filepath string) ([]PackageStatus, error) {
 	file, err := os.Open(filepath)
 	if err != nil {
diff --git a/apt-hook/apthook/types.go b/apt-hook/apthook/types.go
index 9b07694c3f..b4b09ef787 100644
--- a/apt-hook/apthook/types.go
+++ b/apt-hook/apthook/types.go
@@ -46,3 +46,11 @@ type Origin struct {
 	Label    string `json:"label"`
 	Site     string `json:"site"`
 }
+
+type OSRelease struct {
+	Name             string
+	VersionID        string
+	Version          string
+	VersionCodename  string
+	ID               string
+}

From 205667d835f5af9738be5aae378ac32d6c6a1557 Mon Sep 17 00:00:00 2001
From: Dheyay <dheyay.desai@canonical.com>
Date: Fri, 15 Nov 2024 14:03:23 -0800
Subject: [PATCH 3/5] Add debian version comp and main file for go apthook

---
 apt-hook/apthook/chrorder.go   |  21 +++
 apt-hook/apthook/counts.go     | 245 +++++++++++++++++----------------
 apt-hook/apthook/esm_counts.go | 122 +++++++++++++---
 apt-hook/apthook/main.go       |  86 +++++++++---
 apt-hook/apthook/types.go      |  50 +++----
 apt-hook/apthook/version.go    | 135 ++++++++++++++++++
 6 files changed, 474 insertions(+), 185 deletions(-)
 create mode 100644 apt-hook/apthook/chrorder.go
 create mode 100644 apt-hook/apthook/version.go

diff --git a/apt-hook/apthook/chrorder.go b/apt-hook/apthook/chrorder.go
new file mode 100644
index 0000000000..c26e155866
--- /dev/null
+++ b/apt-hook/apthook/chrorder.go
@@ -0,0 +1,21 @@
+// auto-generated, DO NOT EDIT!
+package apthook
+
+var chOrder = [...]int{
+	-5, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271,
+	272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287,
+	288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 299, 300, 301, 302, 303,
+	0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 314, 315, 316, 317, 318, 319,
+	320, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79,
+	80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 347, 348, 349, 350, 351,
+	352, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
+	112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 379, 380, 381, -10, 383,
+	384, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394, 395, 396, 397, 398, 399,
+	400, 401, 402, 403, 404, 405, 406, 407, 408, 409, 410, 411, 412, 413, 414, 415,
+	416, 417, 418, 419, 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431,
+	432, 433, 434, 435, 436, 437, 438, 439, 440, 441, 442, 443, 444, 445, 446, 447,
+	448, 449, 450, 451, 452, 453, 454, 455, 456, 457, 458, 459, 460, 461, 462, 463,
+	464, 465, 466, 467, 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479,
+	480, 481, 482, 483, 484, 485, 486, 487, 488, 489, 490, 491, 492, 493, 494, 495,
+	496, 497, 498, 499, 500, 501, 502, 503, 504, 505, 506, 507, 508, 509, 510, 511,
+}
diff --git a/apt-hook/apthook/counts.go b/apt-hook/apthook/counts.go
index 5c6367f4e5..e4f425f0a2 100644
--- a/apt-hook/apthook/counts.go
+++ b/apt-hook/apthook/counts.go
@@ -1,19 +1,25 @@
 package apthook
 
-import "fmt"
+import (
+	"bufio"
+	"fmt"
+	"os"
+	"strings"
+	"io/ioutil"
+)
 
 const (
-	cloudIDfile = "/run/cloud-init/cloud-id"
-	osReleaseFile = "/etc/os-release"
+	cloudIDFile    = "/run/cloud-init/cloud-id"
+	osReleaseFile  = "/etc/os-release"
 )
 
 type SecurityCounts struct {
 	Standard int
 	ESMInfra int
-	ESMApps int
+	ESMApps  int
 }
 
-// Gets the ubuntu distro for the package
+// getDistroFromPackage retrieves the Ubuntu distribution for the package.
 func getDistroFromPackage(rpc *jsonRPC) string {
 	for _, pkg := range rpc.Params.Packages {
 		for _, origin := range pkg.Versions.Candidate.Origins {
@@ -25,7 +31,7 @@ func getDistroFromPackage(rpc *jsonRPC) string {
 	return ""
 }
 
-// Verify the package is from specific origin
+// verifyOrigin checks if the package is from a specific origin.
 func verifyOrigin(version jsonRPCPackageVersion, origin string) bool {
 	for _, packageOrigin := range version.Origins {
 		if packageOrigin.Origin == origin {
@@ -35,8 +41,8 @@ func verifyOrigin(version jsonRPCPackageVersion, origin string) bool {
 	return false
 }
 
-// Verify the package is from specific origin and archive
-func verifyOriginAndArchive(version jsonRPCPackageVersion, origin string, archive string) bool {
+// verifyOriginAndArchive checks if the package is from a specific origin and archive.
+func verifyOriginAndArchive(version jsonRPCPackageVersion, origin, archive string) bool {
 	for _, packageOrigin := range version.Origins {
 		if packageOrigin.Origin == origin && packageOrigin.Archive == archive {
 			return true
@@ -45,39 +51,19 @@ func verifyOriginAndArchive(version jsonRPCPackageVersion, origin string, archiv
 	return false
 }
 
-// Count security packages from apt stats
-func CountSecurityUpdates(rpc *jsonRPC) *SecurityCounts {
-	counts := &SecurityCounts{}
-	for _, pkg := range rpc.Params.Packages {
-		if pkg.Mode == "upgrade" {
-			// Check esm
-			if verifyOriginAndArchive(pkg.Versions.Install, "UbuntuESMApps", "-apps-security") {
-				counts.ESMApps++
-			}
-			// Checm esm-infra
-			if verifyOriginAndArchive(pkg.Versions.Install, "UbuntuESM", "-infra-security") {
-				counts.ESMInfra++
-			}
-			// Check standard
-			if verifyOriginAndArchive(pkg.Versions.Install, "UbuntuStandard", "-security") {
-				counts.Standard++
-			}
-		}
-	}
-	return counts
-}
-
+// getCloudID retrieves the cloud ID.
 func getCloudID() string {
-	file, err := os.Open(cloudIDfile)
+	file, err := os.Open(cloudIDFile)
 	if err != nil {
 		return ""
 	}
 	defer file.Close()
-	contet, err = ioutil.ReadAll(file)
+
+	content, err := ioutil.ReadAll(file)
 	if err != nil {
 		return ""
 	}
-	fileContent = string(content)
+	fileContent := string(content)
 	if strings.Contains(fileContent, "aws") {
 		return "AWS"
 	} else if strings.Contains(fileContent, "gcp") {
@@ -88,40 +74,37 @@ func getCloudID() string {
 	return ""
 }
 
-// Process os release into a struct
-func parseOSRelease() OSRelease {
+// ParseOSRelease processes OS release information into a struct.
+func ParseOSRelease() OSRelease {
 	file, err := os.Open(osReleaseFile)
 	if err != nil {
-		return ""
+		return OSRelease{}
 	}
 	defer file.Close()
 
-	raw_data := make(map[string]string)
+	rawData := make(map[string]string)
 	scanner := bufio.NewScanner(file)
 	for scanner.Scan() {
-		line := scanner.text()
-		line_split := strings.SplitN(line, "=", 2)
-		if len(line_split) != 2 {
+		line := scanner.Text()
+		lineSplit := strings.SplitN(line, "=", 2)
+		if len(lineSplit) != 2 {
 			continue
 		}
-		key := strings.ToLower(line_split[0])
-		value := strings.Trim(line_split[1], `"`)
-
-		raw_data[key] = value
+		key := strings.ToLower(lineSplit[0])
+		value := strings.Trim(lineSplit[1], `"`)
+		rawData[key] = value
 	}
 
-	osRelease := OSRelease{
-		Name:             raw_data["NAME"],
-		VersionID:        raw_data["VERSION_ID"],
-		Version:          raw_data["VERSION"],
-		VersionCodename:  raw_data["VERSION_CODENAME"],
-		ID:               raw_data["ID"],
+	return OSRelease{
+		Name:            rawData["name"],
+		VersionID:       rawData["version_id"],
+		Version:         rawData["version"],
+		VersionCodename: rawData["version_codename"],
+		ID:              rawData["id"],
 	}
-
-	return osRelease
 }
 
-// Print package names from list
+// printPackageNames prints package names from a list.
 func printPackageNames(packageNames []string) {
 	currLine := " "
 	for _, pkg := range packageNames {
@@ -136,7 +119,72 @@ func printPackageNames(packageNames []string) {
 	}
 }
 
-// CreateCountMessage generates a formatted message describing security updates
+// getESMInfraSeries checks if the OS series is esm_infra.
+func getESMInfraSeries() string {
+	osRelease := ParseOSRelease()
+	switch osRelease.VersionCodename {
+	case "xenial":
+		return "XENIAL"
+	case "bionic":
+		return "BIONIC"
+	default:
+		return "NOT_ESM_INFRA"
+	}
+}
+
+func printLearnMoreContent() {
+	cloudID := getCloudID()
+	esmInfraSeries := getESMInfraSeries()
+
+	switch esmInfraSeries {
+	case "XENIAL":
+		if cloudID == "Azure" {
+			fmt.Println("Learn more about Ubuntu Pro for 16.04 on Azure at https://ubuntu.com/16-04/azure")
+			return
+		}
+		fmt.Println("Learn more about Ubuntu Pro for 16.04 at https://ubuntu.com/16-04")
+		return
+	case "BIONIC":
+		if cloudID == "Azure" {
+			fmt.Println("Learn more about Ubuntu Pro for 18.04 on Azure at https://ubuntu.com/18-04/azure")
+			return
+		}
+		fmt.Println("Learn more about Ubuntu Pro for 18.04 at https://ubuntu.com/18-04")
+		return
+	default:
+		switch cloudID {
+		case "Azure":
+			fmt.Println("Learn more about Ubuntu Pro on Azure at https://ubuntu.com/azure/pro")
+		case "AWS":
+			fmt.Println("Learn more about Ubuntu Pro on AWS at https://ubuntu.com/aws/pro")
+		case "GCP":
+			fmt.Println("Learn more about Ubuntu Pro on GCP at https://ubuntu.com/gcp/pro")
+		default:
+			fmt.Println("Learn more about Ubuntu Pro at https://ubuntu.com/pro")
+		}
+	}
+}
+
+// CountSecurityUpdates counts security packages from apt stats.
+func CountSecurityUpdates(rpc *jsonRPC) *SecurityCounts {
+	counts := &SecurityCounts{}
+	for _, pkg := range rpc.Params.Packages {
+		if pkg.Mode == "upgrade" {
+			if verifyOriginAndArchive(pkg.Versions.Install, "UbuntuESMApps", "-apps-security") {
+				counts.ESMApps++
+			}
+			if verifyOriginAndArchive(pkg.Versions.Install, "UbuntuESM", "-infra-security") {
+				counts.ESMInfra++
+			}
+			if verifyOriginAndArchive(pkg.Versions.Install, "UbuntuStandard", "-security") {
+				counts.Standard++
+			}
+		}
+	}
+	return counts
+}
+
+// GetCountMessage generates a formatted message describing security updates.
 func GetCountMessage(counts *SecurityCounts) string {
 	if counts.Standard == 0 && counts.ESMInfra == 0 && counts.ESMApps == 0 {
 		return ""
@@ -144,16 +192,13 @@ func GetCountMessage(counts *SecurityCounts) string {
 
 	if counts.ESMInfra == 0 && counts.ESMApps == 0 {
 		if counts.Standard == 1 {
-			return fmt.Sprint("1 standard LTS security update")
-		}
-		if counts.Standard > 1 {
-			return fmt.Sprintf("%d standard LTS security updates", counts.Standard)
+			return "1 standard LTS security update"
 		}
+		return fmt.Sprintf("%d standard LTS security updates", counts.Standard)
 	}
 
 	message := make([]string, 0, 3)
 
-	// Add standard count
 	if counts.Standard > 0 {
 		if counts.Standard == 1 {
 			message = append(message, "1 standard LTS security update")
@@ -162,7 +207,6 @@ func GetCountMessage(counts *SecurityCounts) string {
 		}
 	}
 
-	// Add ESM Infra count
 	if counts.ESMInfra > 0 {
 		if counts.ESMInfra == 1 {
 			message = append(message, "1 esm-infra security update")
@@ -171,7 +215,6 @@ func GetCountMessage(counts *SecurityCounts) string {
 		}
 	}
 
-	// Add ESM Apps count
 	if counts.ESMApps > 0 {
 		if counts.ESMApps == 1 {
 			message = append(message, "1 esm-apps security update")
@@ -181,18 +224,18 @@ func GetCountMessage(counts *SecurityCounts) string {
 	}
 
 	switch len(message) {
-		case 1:
-			return message[0]
-		case 2:
-			return fmt.Sprintf("%s and %s", message[0], message[1])
-		case 3:
-			return fmt.Sprintf("%s, %s and %s", message[0], message[1], message[2])
+	case 1:
+		return message[0]
+	case 2:
+		return fmt.Sprintf("%s and %s", message[0], message[1])
+	case 3:
+		return fmt.Sprintf("%s, %s and %s", message[0], message[1], message[2])
 	}
 
 	return ""
 }
 
-
+// CollectProPackagesFromRPC collects packages from RPC.
 func CollectProPackagesFromRPC(rpc *jsonRPC) []string {
 	var expiredPackages []string
 
@@ -223,7 +266,7 @@ func CollectProPackagesFromRPC(rpc *jsonRPC) []string {
 	return expiredPackages
 }
 
-// PrintExpiredProPackages logs messages about expired pro package status
+// PrintExpiredProPackages logs messages about expired pro package status.
 func PrintExpiredProPackages(packageNames []string) {
 	if len(packageNames) == 0 {
 		return
@@ -233,64 +276,22 @@ func PrintExpiredProPackages(packageNames []string) {
 	fmt.Println("Renew your subscription or run `sudo pro detach` to remove these errors")
 }
 
-// Seperate function to check if os series is esm_infra
-func getESMInfraSeries() string {
-	osRelease := parseOSRelease()
-	if osRelease.VersionCodename == "xenial" {
-		return "XENIAL"
-	} else if osRelease.VersionCodename == "bionic" {
-		return "BIONIC"
-	}
-	return "NOT_ESM_INFRA"
-}
-
-func printLearnMoreContent() {
-	cloudID := getCloudID()
-	esmInfraSeries := getOSRelease()
-
-	if esmInfraSeries == "XENIAL" {
-		if cloudID == "Azure" {
-			fmt.Printf("Learn more about Ubuntu Pro for 16.04 on Azure at %s\n",
-				"https://ubuntu.com/16-04/azure")
-			return
+// printESMPackages prints expired ESM packages.
+func PrintESMPackages(esmType string, packageNames []string) {
+	if esmType == "APPS" {
+		if len(packageNames) == 1 {
+			fmt.Println("Get another security update through Ubuntu Pro with 'esm-apps' enabled:")
 		} else {
-			fmt.Printf("Learn more about Ubuntu Pro for 16.04 at %s\n",
-				"https://ubuntu.com/16-04")
-			return
+			fmt.Println("Get more security updates through Ubuntu Pro with 'esm-apps' enabled:")
 		}
-	} else if esmInfraSeries == "BIONIC" {
-		if cloudID == "Azure" {
-			fmt.Printf("Learn more about Ubuntu Pro for 18.04 on Azure at %s\n",
-				"https://ubuntu.com/18-04/azure")
-			return
+	} else if esmType == "INFRA" {
+		if len(packageNames) == 1 {
+			fmt.Println("The following security update requires Ubuntu Pro with 'esm-infra' enabled:")
 		} else {
-			fmt.Printf("Learn more about Ubuntu Pro for 18.04 at %s\n",
-				"https://ubuntu.com/18-04")
-			return
-		}
-	} else {
-		if cloudID == "Azure" {
-			fmt.Printf("Learn more about Ubuntu Pro on Azure at %s\n",
-				"https://ubuntu.com/azure/pro")
-			return
-		} else if cloudID == "AWS" {
-			fmt.Printf("Learn more about Ubuntu Pro on AWS at %s\n",
-				"https://ubuntu.com/aws/pro")
-			return
-		} else if cloudID == "GCP" {
-			fmt.Printf("Learn more about Ubuntu Pro on GCP at %s\n",
-				"https://ubuntu.com/gcp/pro")
-			return
+			fmt.Println("The following security updates require Ubuntu Pro with 'esm-infra' enabled:")
 		}
 	}
 
-	fmt.Printf("Learn more about Ubuntu Pro at %s\n",
-		"https://ubuntu.com/pro")
-	return
-}
-
-func printESMPackages(esmType ESMType, packageNames []string) {
-	fmt.Println("The following packages will fail to download because your Ubuntu Pro subscription has expired")
 	printPackageNames(packageNames)
-	fmt.Println("Renew your subscription or `sudo pro detach` to remove these errors")
+	printLearnMoreContent()
 }
diff --git a/apt-hook/apthook/esm_counts.go b/apt-hook/apthook/esm_counts.go
index ec583f8315..c2acbda9d8 100644
--- a/apt-hook/apthook/esm_counts.go
+++ b/apt-hook/apthook/esm_counts.go
@@ -2,42 +2,65 @@ package apthook
 
 import (
 	"fmt"
+	"path/filepath"
 	"strings"
+	"bufio"
+	"os"
+	"os/exec"
 )
 
 const (
-	systemPkgStatus = "/var/lib/dpkg/status"
-	esmPkgStatus	= "/var/lib/ubuntu-advantage/apt-esm/var/lib/dpkg/status"
+	systemPkgPath = "/var/lib/dpkg/status"
 )
-// Instead of esmPkgStatus
-// we need to use the following from var/lib/ubuntu-advantage/apt-esm/var/lib/apt/lists:
-// esmAppsSecurityList
-// esmAppsUpdatesList
-// esmInfraSecurityList
-// esmInfraUpdatesList
 
-
-type PackageStatus struct {
+type PkgStatus struct {
 	Name	string
 	Version string
 	Status  string
 	Source  string
 }
 
-// Func to get file pathes for esm packages
-// Based on os release codename
+type ESMUpdates struct {
+	InfraPackages []string
+	AppsPackages  []string
+}
+
+func getDpkgArch() string {
+	cmd := exec.Command("dpkg", "--print-architecture")
+	out, err := cmd.Output()
+	if err != nil {
+		return ""
+	}
+	arch := strings.TrimSpace(string(out))
+	return arch
+}
+
+func getEsmPackagesFilePath(osRelease OSRelease) (string, string) {
+    curSeries := osRelease.VersionCodename
+    curArch := getDpkgArch()
+
+    const baseDir = "/var/lib/ubuntu-advantage/apt-esm/var/lib/apt/lists"
+
+    security := filepath.Join(baseDir,
+        fmt.Sprintf("esm.ubuntu.com_apps_ubuntu_dists_%s-apps-security_main_binary-%s_Packages",
+            curSeries, curArch))
 
-// Packages need to be read and parsed from all above files
-// and then returned as a slice of PackageStatus
-func readPackages(filepath string) ([]PackageStatus, error) {
+    updates := filepath.Join(baseDir,
+        fmt.Sprintf("esm.ubuntu.com_apps_ubuntu_dists_%s-apps-updates_main_binary-%s_Packages",
+            curSeries, curArch))
+
+    return security, updates
+}
+
+func readPackages(filepath string) ([]PkgStatus, error) {
 	file, err := os.Open(filepath)
 	if err != nil {
-		return nil, error
+		return nil, err
 	}
 	defer file.Close()
 
-	var packages []PackageStatus
-	var curPackage PackageStatus
+	var packages []PkgStatus
+	var curPackage PkgStatus
 
 	scanner := bufio.NewScanner(file)
 	parsingPackage := false
@@ -47,10 +70,10 @@ func readPackages(filepath string) ([]PackageStatus, error) {
 
 		if strings.HasPrefix(line, "Package:") {
 			if parsingPackage {
-				packages = append(packages, currentPackage)
+				packages = append(packages, curPackage)
 			}
 
-			curPackage = PackageStatus{}
+			curPackage = PkgStatus{}
 			curPackage.Name = strings.TrimSpace(strings.TrimPrefix(line, "Package:"))
 			parsingPackage = true
 		} else if strings.HasPrefix(line, "Version:") {
@@ -70,5 +93,64 @@ func readPackages(filepath string) ([]PackageStatus, error) {
 		packages = append(packages, curPackage)
 	}
 
+	// Can return a dict (map) for faster lookup instead of iterating?
 	return packages, nil
 }
+
+// GetPotentialESMUpdates returns a list of packages that are available for
+// updates in the ESM repo
+func GetPotentialESMUpdates() (*ESMUpdates, error) {
+	updates := &ESMUpdates{}
+
+	systemPkgs, err := readPackages(systemPkgPath)
+	if err != nil {
+		return nil, fmt.Errorf("reading system packages: %w", err)
+	}
+
+	// Get ESM package list paths (currently only apps)
+	osRelease := ParseOSRelease()
+	securityPath, updatesPath := getEsmPackagesFilePath(osRelease)
+
+	securityPkgs, err := readPackages(securityPath)
+	if err != nil {
+		return nil, fmt.Errorf("reading security packages: %w", err)
+	}
+
+	updatesPkgs, err := readPackages(updatesPath)
+	if err != nil {
+		return nil, fmt.Errorf("reading updates packages: %w", err)
+	}
+
+	for _, sysPkg := range systemPkgs {
+		// If installed
+		if !strings.Contains(sysPkg.Status, "installed") {
+			continue
+		}
+
+		// Check security packages path
+		for _, secPkg := range securityPkgs {
+			if secPkg.Name == sysPkg.Name && CompareVersions(secPkg.Version, sysPkg.Version) < 0{
+				// Check ESM Source
+				if strings.Contains(secPkg.Source, "UbuntuESMApps") {
+					updates.AppsPackages = append(updates.AppsPackages, secPkg.Name)
+				} else if strings.Contains(secPkg.Source, "UbuntuESM") {
+					updates.InfraPackages = append(updates.InfraPackages, secPkg.Name)
+				}
+			}
+		}
+
+		// Check updates packages path
+		for _, updatesPkg := range updatesPkgs {
+			if updatesPkg.Name == sysPkg.Name && CompareVersions(updatesPkg.Version, sysPkg.Version) < 0{
+				// Check ESM Source
+				if strings.Contains(updatesPkg.Source, "UbuntuESMApps") {
+					updates.AppsPackages = append(updates.AppsPackages, updatesPkg.Name)
+				} else if strings.Contains(updatesPkg.Source, "UbuntuESM") {
+					updates.InfraPackages = append(updates.InfraPackages, updatesPkg.Name)
+				}
+			}
+		}
+	}
+
+	return updates, nil
+}
diff --git a/apt-hook/apthook/main.go b/apt-hook/apthook/main.go
index 2b9c239966..37373025dc 100644
--- a/apt-hook/apthook/main.go
+++ b/apt-hook/apthook/main.go
@@ -2,9 +2,8 @@ package apthook
 
 import (
 	"fmt"
-	"strings"
 	"os"
-	"os/exec"
+	"strconv"
 )
 
 const (
@@ -14,20 +13,71 @@ const (
 	expiredNotice = "/var/lib/ubuntu-advantage/notices/5-contract_expired"
 )
 
-func main() {
-	// Get apt hook socker
-	// Make sure socket in not empty and exists
-	// jsonrpc handshake
-	// jsonrpc read rpc message
-	//
-	// if method status
-	// 	display security count message
-	// if method pre-prompt
-	//  get potential esm updates
-	//  call display functions from counts file
-	//
-	// Display apt news
-	// Display contract expiration notice
-	//
-	// jsonpc bye bye
+func main() error {
+	sockFd := os.Getenv("APT_HOOK_SOCKET")
+	if sockFd == "" {
+		fmt.Errorf("pro-apt-hook: missing socket fd")
+	}
+
+	fd, err := strconv.Atoi(sockFd)
+	if err != nil {
+		fmt.Errorf("pro-apt-hook: invalid socket fd: %w", err)
+	}
+
+	file := os.NewFile(uintptr(fd), "apt-hook-socket")
+	if file == nil {
+		fmt.Errorf("pro-apt-hook: cannot open file descriptor %d", fd)
+	}
+	defer file.Close()
+
+	conn, err := NewConnection(file)
+	if err != nil {
+		fmt.Errorf("pro-apt-hook: failed to create connection: %w", err)
+	}
+	defer conn.Close()
+
+	if err := conn.Handshake(); err != nil {
+		fmt.Errorf("pro-apt-hook: handshake failed: %w", err)
+	}
+
+	msg, err := conn.ReadMessage()
+	if err != nil {
+		fmt.Errorf("pro-apt-hook: reading message: %w", err)
+	}
+
+	switch msg.Method {
+
+	case methodStats:
+		counts := CountSecurityUpdates(msg)
+		if message := GetCountMessage(counts); message != "" {
+			fmt.Println(message)
+		}
+
+	case methodPrePrompt:
+		// Get ESM updates
+		if updates, err := GetPotentialESMUpdates(); err == nil {
+			if len(updates.InfraPackages) > 0 {
+				PrintESMPackages("INFRA", updates.InfraPackages)
+			} else if len(updates.AppsPackages) > 0 {
+				PrintESMPackages("APPS", updates.AppsPackages)
+			}
+		}
+
+		if news, err := os.ReadFile(aptNewsFile); err == nil {
+			fmt.Print(string(news))
+		}
+
+		if _, err := os.Stat(expiredNotice); err == nil {
+			expiredPkgs := CollectProPackagesFromRPC(msg)
+			if len(expiredPkgs) > 0 {
+				PrintExpiredProPackages(expiredPkgs)
+			}
+		}
+	}
+
+	if err := conn.Bye(); err != nil {
+		fmt.Errorf("pro-apt-hook: bye failed: %w", err)
+	}
+
+	return nil
 }
diff --git a/apt-hook/apthook/types.go b/apt-hook/apthook/types.go
index b4b09ef787..4be9827ec6 100644
--- a/apt-hook/apthook/types.go
+++ b/apt-hook/apthook/types.go
@@ -5,36 +5,36 @@ import (
 )
 
 type jsonRPC struct {
-	JsonRPC	string	`json:"jsonrpc"`
-	Method	string	`json:"method"`
-	Params	struct {
-		Command	        string   `json:"command"`
-		UnknownPackages	[]string `json:"unknown-packages"`
-		Packages		[]Package `json:"packages"`
+	JsonRPC string `json:"jsonrpc"`
+	Method  string `json:"method"`
+	Params  struct {
+		Command         string    `json:"command"`
+		UnknownPackages []string  `json:"unknown-packages"`
+		Packages        []Package `json:"packages"`
 	} `json:"params"`
 }
 
 type Package struct {
-	ID           int    `json:"id"`
-	Name         string `json:"name"`
-	Architecture string `json:"architecture"`
-	Mode         string `json:"mode"`
-	Versions     Versions `json:"versions,omitempty"` // optional
-	Automatic    bool `json:"automatic,omitempty"`	// optional
-	Current	     string `json:"current,omitempty"`	// optional
+	ID           int      `json:"id"`
+	Name         string   `json:"name"`
+	Architecture string   `json:"architecture"`
+	Mode         string   `json:"mode"`
+	Versions     Versions `json:"versions,omitempty"`  // Optional
+	Automatic    bool     `json:"automatic,omitempty"` // Optional
+	Current      string   `json:"current,omitempty"`   // Optional
 }
 
 type Versions struct {
-	Candidate	jsonRPCPackageVersion	`json:"candidate"`
-	Install		jsonRPCPPackageVersion	`json:"install"`
-	Current		jsonRPCPackageVersion	`json:"current"`
+	Candidate jsonRPCPackageVersion `json:"candidate"`
+	Install   jsonRPCPackageVersion  `json:"install"`
+	Current   jsonRPCPackageVersion  `json:"current"`
 }
 
 type jsonRPCPackageVersion struct {
-	ID           int    `json:"id"`
-	Version      string `json:"version"`
-	Architecture string `json:"architecture"`
-	Pin          int    `json:"pin"`
+	ID           int      `json:"id"`
+	Version      string   `json:"version"`
+	Architecture string   `json:"architecture"`
+	Pin          int      `json:"pin"`
 	Origins      []Origin `json:"origins"`
 }
 
@@ -48,9 +48,9 @@ type Origin struct {
 }
 
 type OSRelease struct {
-	Name             string
-	VersionID        string
-	Version          string
-	VersionCodename  string
-	ID               string
+	Name            string `json:"name"`
+	VersionID       string `json:"version_id"`
+	Version         string `json:"version"`
+	VersionCodename string `json:"version_codename"`
+	ID              string `json:"id"`
 }
diff --git a/apt-hook/apthook/version.go b/apt-hook/apthook/version.go
new file mode 100644
index 0000000000..0fd2e93f13
--- /dev/null
+++ b/apt-hook/apthook/version.go
@@ -0,0 +1,135 @@
+// From the chisel library for debian  version comparison
+
+package apthook
+
+import (
+	"strings"
+)
+
+func max(a, b int) int {
+	if a < b {
+		return b
+	}
+	return a
+}
+
+//go:generate go run ./chrorder/main.go -package=deb -output=chrorder.go
+
+func cmpString(as, bs string) int {
+	for i := 0; i < max(len(as), len(bs)); i++ {
+		var a uint8
+		var b uint8
+		if i < len(as) {
+			a = as[i]
+		}
+		if i < len(bs) {
+			b = bs[i]
+		}
+		if chOrder[a] < chOrder[b] {
+			return -1
+		}
+		if chOrder[a] > chOrder[b] {
+			return +1
+		}
+	}
+	return 0
+}
+
+func trimLeadingZeroes(a string) string {
+	for i := 0; i < len(a); i++ {
+		if a[i] != '0' {
+			return a[i:]
+		}
+	}
+	return ""
+}
+
+// a and b both match /[0-9]+/
+func cmpNumeric(a, b string) int {
+	a = trimLeadingZeroes(a)
+	b = trimLeadingZeroes(b)
+
+	switch d := len(a) - len(b); {
+	case d > 0:
+		return 1
+	case d < 0:
+		return -1
+	}
+	for i := 0; i < len(a); i++ {
+		switch {
+		case a[i] > b[i]:
+			return 1
+		case a[i] < b[i]:
+			return -1
+		}
+	}
+	return 0
+}
+
+func nextFrag(s string) (frag, rest string, numeric bool) {
+	if len(s) == 0 {
+		return "", "", false
+	}
+
+	var i int
+	if s[0] >= '0' && s[0] <= '9' {
+		// is digit
+		for i = 1; i < len(s) && s[i] >= '0' && s[i] <= '9'; i++ {
+		}
+		numeric = true
+	} else {
+		// not digit
+		for i = 1; i < len(s) && (s[i] < '0' || s[i] > '9'); i++ {
+		}
+	}
+	return s[:i], s[i:], numeric
+}
+
+func compareSubversion(va, vb string) int {
+	var a, b string
+	var anum, bnum bool
+	var res int
+	for res == 0 {
+		a, va, anum = nextFrag(va)
+		b, vb, bnum = nextFrag(vb)
+		if a == "" && b == "" {
+			break
+		}
+		if anum && bnum {
+			res = cmpNumeric(a, b)
+		} else {
+			res = cmpString(a, b)
+		}
+	}
+	return res
+}
+
+// CompareVersions compare two version strings that follow the debian
+// version policy and
+// Returns:
+//
+//	-1 if a is smaller than b
+//	 0 if a equals b
+//	+1 if a is bigger than b
+func CompareVersions(va, vb string) int {
+	var sa, sb string
+	if ia := strings.IndexByte(va, '-'); ia < 0 {
+		sa = "0"
+	} else {
+		va, sa = va[:ia], va[ia+1:]
+	}
+	if ib := strings.IndexByte(vb, '-'); ib < 0 {
+		sb = "0"
+	} else {
+		vb, sb = vb[:ib], vb[ib+1:]
+	}
+
+	// the main version number (before the "-")
+	res := compareSubversion(va, vb)
+	if res != 0 {
+		return res
+	}
+
+	// the subversion revision behind the "-"
+	return compareSubversion(sa, sb)
+}

From 4ebe79f433706ddce6b707853cd84e91b2a3e4c8 Mon Sep 17 00:00:00 2001
From: Dheyay <dheyay.desai@canonical.com>
Date: Fri, 15 Nov 2024 14:10:56 -0800
Subject: [PATCH 4/5] add go mod file

---
 apt-hook/apthook/go.mod | 3 +++
 1 file changed, 3 insertions(+)
 create mode 100644 apt-hook/apthook/go.mod

diff --git a/apt-hook/apthook/go.mod b/apt-hook/apthook/go.mod
new file mode 100644
index 0000000000..6cf1184a70
--- /dev/null
+++ b/apt-hook/apthook/go.mod
@@ -0,0 +1,3 @@
+module jsonhook
+
+go 1.16

From 51a95890feeb92deccd81d23ff2b5ccd8ddefb84 Mon Sep 17 00:00:00 2001
From: Dheyay <dheyay.desai@canonical.com>
Date: Fri, 15 Nov 2024 14:36:33 -0800
Subject: [PATCH 5/5] go linter

---
 apt-hook/apthook/counts.go | 12 ------------
 apt-hook/apthook/main.go   | 14 +++++++-------
 apt-hook/apthook/rpc.go    |  9 +++++----
 apt-hook/apthook/types.go  |  4 ----
 4 files changed, 12 insertions(+), 27 deletions(-)

diff --git a/apt-hook/apthook/counts.go b/apt-hook/apthook/counts.go
index e4f425f0a2..b99a17ba89 100644
--- a/apt-hook/apthook/counts.go
+++ b/apt-hook/apthook/counts.go
@@ -19,18 +19,6 @@ type SecurityCounts struct {
 	ESMApps  int
 }
 
-// getDistroFromPackage retrieves the Ubuntu distribution for the package.
-func getDistroFromPackage(rpc *jsonRPC) string {
-	for _, pkg := range rpc.Params.Packages {
-		for _, origin := range pkg.Versions.Candidate.Origins {
-			if origin.Codename != "" {
-				return origin.Codename
-			}
-		}
-	}
-	return ""
-}
-
 // verifyOrigin checks if the package is from a specific origin.
 func verifyOrigin(version jsonRPCPackageVersion, origin string) bool {
 	for _, packageOrigin := range version.Origins {
diff --git a/apt-hook/apthook/main.go b/apt-hook/apthook/main.go
index 37373025dc..ff275192f6 100644
--- a/apt-hook/apthook/main.go
+++ b/apt-hook/apthook/main.go
@@ -16,33 +16,33 @@ const (
 func main() error {
 	sockFd := os.Getenv("APT_HOOK_SOCKET")
 	if sockFd == "" {
-		fmt.Errorf("pro-apt-hook: missing socket fd")
+		return fmt.Errorf("pro-apt-hook: missing socket fd")
 	}
 
 	fd, err := strconv.Atoi(sockFd)
 	if err != nil {
-		fmt.Errorf("pro-apt-hook: invalid socket fd: %w", err)
+		return fmt.Errorf("pro-apt-hook: invalid socket fd: %w", err)
 	}
 
 	file := os.NewFile(uintptr(fd), "apt-hook-socket")
 	if file == nil {
-		fmt.Errorf("pro-apt-hook: cannot open file descriptor %d", fd)
+		return fmt.Errorf("pro-apt-hook: cannot open file descriptor %d", fd)
 	}
 	defer file.Close()
 
 	conn, err := NewConnection(file)
 	if err != nil {
-		fmt.Errorf("pro-apt-hook: failed to create connection: %w", err)
+		return fmt.Errorf("pro-apt-hook: failed to create connection: %w", err)
 	}
 	defer conn.Close()
 
 	if err := conn.Handshake(); err != nil {
-		fmt.Errorf("pro-apt-hook: handshake failed: %w", err)
+		return fmt.Errorf("pro-apt-hook: handshake failed: %w", err)
 	}
 
 	msg, err := conn.ReadMessage()
 	if err != nil {
-		fmt.Errorf("pro-apt-hook: reading message: %w", err)
+		return fmt.Errorf("pro-apt-hook: reading message: %w", err)
 	}
 
 	switch msg.Method {
@@ -76,7 +76,7 @@ func main() error {
 	}
 
 	if err := conn.Bye(); err != nil {
-		fmt.Errorf("pro-apt-hook: bye failed: %w", err)
+		return fmt.Errorf("pro-apt-hook: bye failed: %w", err)
 	}
 
 	return nil
diff --git a/apt-hook/apthook/rpc.go b/apt-hook/apthook/rpc.go
index 8efa87367a..b1ff490a29 100644
--- a/apt-hook/apthook/rpc.go
+++ b/apt-hook/apthook/rpc.go
@@ -1,11 +1,12 @@
 package apthook
 
-importy (
+import (
 	"bufio"
 	"encoding/json"
 	"io"
 	"net"
 	"fmt"
+	"os"
 )
 
 const (
@@ -27,7 +28,7 @@ func (c *Connection) ReadMessage() (*jsonRPC, error) {
 }
 
 func (c *Connection) WriteResponse(version string, id int) error {
-	respo := fmt.Sprintf(`{"jsonrpc": "2.0", "id": %d, "result": {"version": "%s"}}`, id, version)
+	response := fmt.Sprintf(`{"jsonrpc": "2.0", "id": %d, "result": {"version": "%s"}}`, id, version)
 	_, err := c.conn.Write([]byte(response + "\n\n"))
 	return err
 }
@@ -81,13 +82,13 @@ func readRPC(r *bufio.Reader) (*jsonRPC, error) {
 
 	var msg jsonRPC
 	if err := json.Unmarshal(line, &msg); err != nil {
-		return nul, fmt.Errorf("parsing json: %w", err)
+		return nil, fmt.Errorf("parsing json: %w", err)
 	}
 
 	emptyLine, _, err := r.ReadLine()
 
 	if err != nil {
-		return nul, fmt.Errorf("reading empty line: %w", err)
+		return nil, fmt.Errorf("reading empty line: %w", err)
 	}
 
 	if string(emptyLine) != "" {
diff --git a/apt-hook/apthook/types.go b/apt-hook/apthook/types.go
index 4be9827ec6..8c33c4dc2f 100644
--- a/apt-hook/apthook/types.go
+++ b/apt-hook/apthook/types.go
@@ -1,9 +1,5 @@
 package apthook
 
-import (
-	"encoding/json"
-)
-
 type jsonRPC struct {
 	JsonRPC string `json:"jsonrpc"`
 	Method  string `json:"method"`