path: root/src/main.rs

use std::process::Command;
use std::str::FromStr;
use std::{fs, io};

use tonic::{Request, Response, Status, transport::Server};

use crate::devinfod::dev_info_server::{DevInfo, DevInfoServer};
use crate::devinfod::{
    GetBuildInfoRequest, GetBuildInfoResponse, GetRootRequest, GetRootResponse,
    GetSwVersionsRequest, GetSwVersionsResponse, GetUbootEnvRequest, GetUbootEnvResponse,
};

pub mod devinfod {
    tonic::include_proto!("unit.info.v0");
}

#[derive(Debug, Default)]
pub struct MyDevInfo {}

pub struct Distro {
    name: String,
    version: String,
}

impl Distro {
    pub fn to_pb(&self) -> devinfod::Distro {
        devinfod::Distro {
            name: self.name.clone(),
            version: self.version.clone(),
        }
    }
}

pub struct Layer {
    name: String,
    branch: String,
    revision: String,
}

impl Layer {
    pub fn to_pb(&self) -> devinfod::Layer {
        devinfod::Layer {
            name: self.name.clone(),
            branch: self.branch.clone(),
            revision: self.revision.clone(),
        }
    }
}

enum BuildInfoSection {
    BuildConfiguration,
    LayerRevisions,
}

impl BuildInfoSection {
    pub fn from_line(line: &str) -> Option<Self> {
        if line.starts_with("Build Configuration") {
            Some(BuildInfoSection::BuildConfiguration)
        } else if line.starts_with("Layer Revisions") {
            Some(BuildInfoSection::LayerRevisions)
        } else {
            None
        }
    }
}

pub struct BuildInfo {
    build_configuration: Distro,
    layer_revisions: Vec<Layer>,
}

impl BuildInfo {
    pub fn to_pb(&self) -> devinfod::BuildInfo {
        let mut layers_pb = Vec::<devinfod::Layer>::new();

        for layer in &self.layer_revisions {
            layers_pb.push(layer.to_pb());
        }

        devinfod::BuildInfo {
            distro: Some(self.build_configuration.to_pb()),
            layers: layers_pb,
        }
    }
}

impl FromStr for BuildInfo {
    type Err = io::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut distro_name: Option<String> = None;
        let mut distro_version: Option<String> = None;
        let mut layers = Vec::new();
        let mut curr_section: Option<BuildInfoSection> = None;

        for line in s.lines() {
            let trimmed = line.trim();

            if line.ends_with("|") {
                if let Some(section) = BuildInfoSection::from_line(line) {
                    curr_section = Some(section);
                    continue;
                }
            }

            if line.starts_with("-----") {
                continue;
            }

            if let Some(section) = &curr_section {
                match section {
                    BuildInfoSection::BuildConfiguration => {
                        if let Some((key, value)) = trimmed.split_once("=") {
                            let key = key.trim();
                            let value = value.trim();

                            match key {
                                "DISTRO" => distro_name = Some(value.to_string()),
                                "DISTRO_VERSION" => distro_version = Some(value.to_string()),
                                _ => {
                                    return Err(io::Error::new(
                                        io::ErrorKind::InvalidData,
                                        format!(
                                            "Distro configuration contains unexpected key: {}",
                                            key
                                        ),
                                    ));
                                }
                            }
                        } else {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                "Entry for build configuration has invalid format, should be of form <key> = <value>",
                            ));
                        }
                    }
                    BuildInfoSection::LayerRevisions => {
                        if let Some((key, value)) = trimmed.split_once("=") {
                            let layer_name = key.trim();
                            let value = value.trim();
                            if let Some((branch, revision)) = value.split_once(":") {
                                layers.push(Layer {
                                    name: layer_name.to_string(),
                                    branch: branch.to_string(),
                                    revision: revision.to_string(),
                                });
                            } else {
                                return Err(io::Error::new(
                                    io::ErrorKind::InvalidData,
                                    format!(
                                        "Entry for layer {} contains invalid value format, should be of form <branch>:<revision>",
                                        layer_name
                                    ),
                                ));
                            }
                        } else {
                            return Err(io::Error::new(
                                io::ErrorKind::InvalidData,
                                "Entry for layer revisions has invalid format, should be of form <key> = <value>",
                            ));
                        }
                    }
                }
            }
        }

        match (distro_name, distro_version) {
            (Some(name), Some(version)) => Ok(BuildInfo {
                build_configuration: Distro { name, version },
                layer_revisions: layers,
            }),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "Missing required build configuration information",
            )),
        }
    }
}

pub struct SwVersion {
    name: String,
    version: String,
}

impl SwVersion {
    pub fn to_pb(&self) -> devinfod::SwVersion {
        devinfod::SwVersion {
            name: self.name.clone(),
            version: self.version.clone(),
        }
    }
}

pub struct SwVersions {
    versions: Vec<SwVersion>,
}

impl SwVersions {
    pub fn to_pb(&self) -> devinfod::SwVersions {
        let mut versions = Vec::<devinfod::SwVersion>::new();

        for sw_version in &self.versions {
            versions.push(sw_version.to_pb());
        }

        devinfod::SwVersions { versions }
    }
}

impl FromStr for SwVersions {
    type Err = io::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut versions = Vec::new();

        for line in s.lines() {
            let trimmed = line.trim();

            if let Some((key, value)) = trimmed.split_once(" ") {
                let key = key.trim();
                let value = value.trim();

                versions.push(SwVersion {
                    name: key.to_string(),
                    version: value.to_string(),
                });
            } else {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "Entry for sw versions has invalid format, should be of form <key> <value>",
                ));
            }
        }

        Ok(SwVersions { versions })
    }
}

fn swupdate_get_root() -> io::Result<String> {
    let output = Command::new("swupdate").arg("-g").output()?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        eprintln!("swupdate stderr: {}", stderr);
        return Err(io::Error::new(
            io::ErrorKind::Other,
            format!("swupdate failed with status: {}", output.status),
        ));
    }

    let root = String::from_utf8(output.stdout)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
        .trim()
        .to_string();

    Ok(root)
}

struct EnvVar {
    name: String,
    value: String,
}

impl EnvVar {
    pub fn to_pb(&self) -> devinfod::EnvVar {
        devinfod::EnvVar {
            name: self.name.clone(),
            value: self.value.clone(),
        }
    }
}

struct UbootEnv {
    vars: Vec<EnvVar>,
}

impl UbootEnv {
    pub fn from_cmd() -> io::Result<Self> {
        let output = Command::new("fw_printenv").output()?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            eprintln!("fw_printenv stderr: {}", stderr);
            return Err(io::Error::new(
                io::ErrorKind::Other,
                format!("fw_printenv failed with status: {}", output.status),
            ));
        }

        let output_str = String::from_utf8(output.stdout)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
            .trim()
            .to_string();

        let mut vars = Vec::<EnvVar>::new();

        for line in output_str.lines() {
            let trimmed = line.trim();
            if let Some((key, value)) = trimmed.split_once("=") {
                let key = key.trim();

                if matches!(
                    key,
                    "target_partition_device_node" | "upgrade_available" | "ustate"
                ) {
                    let value = value.trim();
                    vars.push(EnvVar {
                        name: key.to_string(),
                        value: value.to_string(),
                    })
                }
            } else {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "Line for u-boot environment has invalid format, should be of form <key>=<value>",
                ));
            }
        }

        Ok(UbootEnv { vars })
    }

    pub fn to_pb(&self) -> devinfod::UbootEnv {
        let mut vars_pb = Vec::<devinfod::EnvVar>::new();

        for var in &self.vars {
            vars_pb.push(var.to_pb());
        }

        devinfod::UbootEnv { vars: vars_pb }
    }
}

#[tonic::async_trait]
impl DevInfo for MyDevInfo {
    async fn get_build_info(
        &self,
        _request: Request<GetBuildInfoRequest>,
    ) -> Result<Response<GetBuildInfoResponse>, Status> {
        let build_info: BuildInfo = fs::read_to_string("/etc/buildinfo")?.parse()?;
        let build_info_pb = build_info.to_pb();
        let response = GetBuildInfoResponse {
            info: Some(build_info_pb),
        };
        Ok(Response::new(response))
    }

    async fn get_root(
        &self,
        _request: Request<GetRootRequest>,
    ) -> Result<Response<GetRootResponse>, Status> {
        let root = swupdate_get_root()?;
        let response = GetRootResponse { root };
        Ok(Response::new(response))
    }

    async fn get_sw_versions(
        &self,
        _request: Request<GetSwVersionsRequest>,
    ) -> Result<Response<GetSwVersionsResponse>, Status> {
        let sw_versions: SwVersions = fs::read_to_string("/etc/sw-versions")?.parse()?;
        let sw_versions_pb = sw_versions.to_pb();
        let response = GetSwVersionsResponse {
            versions: Some(sw_versions_pb),
        };
        Ok(Response::new(response))
    }

    async fn get_uboot_env(
        &self,
        _request: Request<GetUbootEnvRequest>,
    ) -> Result<Response<GetUbootEnvResponse>, Status> {
        let env = UbootEnv::from_cmd()?;
        let env_pb = env.to_pb();
        let response = GetUbootEnvResponse { env: Some(env_pb) };
        Ok(Response::new(response))
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let addr = "0.0.0.0:60067".parse().unwrap();
    let devinfo = MyDevInfo::default();

    println!("Listening on {}", addr);

    Server::builder()
        .add_service(DevInfoServer::new(devinfo))
        .serve(addr)
        .await?;

    Ok(())
}