testctl/src/main.rs

use core::{str, task};
use std::collections::HashSet;

use anyhow::{anyhow, Result};
use clap::{Parser, Subcommand, ValueEnum};
use env_logger::Target;
use futures_lite::future::block_on;
use log::{debug, error, info};
use nusb::{
    transfer::{ControlIn, ControlOut, ControlType, Recipient},
    Interface,
};
use serde::{Deserialize, Serialize};

#[derive(Parser)]
#[command(version, about, long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Clone, Copy, Debug, ValueEnum)]
enum SwdTarget {
    Dut,
    Buddy,
}

#[derive(Clone, Subcommand)]
enum Commands {
    /// Query test target for info.
    Info,

    /// Sets the Test Status RGB LED color.
    SetLed { red: u8, green: u8, blue: u8 },
    /// Selects which target is connected to the probe.
    Target { target: SwdTarget },
}

#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
enum TestTargetCapability {
    Uart,
    Spi,
    I2c,
    DigitalIo,
    Analog,
}

#[derive(Debug, Deserialize, Serialize)]
struct TestTargetInfo {
    /// Name of the test target.
    name: String,

    /// MCU of the device under test.
    mcu: String,

    /// Capabilities of the test target board.
    capabilities: Vec<TestTargetCapability>,
}

const RPI_VENDOR_ID: u16 = 0x2e8a;
const DEBUG_PROBE_PRODUCT_ID: u16 = 0x000c;

struct PigweedTestTarget {
    interface: Interface,
}

#[derive(Debug)]
#[repr(u16)]
enum PigweedVendorCommand {
    Hello = 0,
    SetTarget = 1,
    GetTargetInfo = 2,
    SetLed = 3,
}

impl PigweedTestTarget {
    pub fn new() -> Result<Self> {
        let device_info = nusb::list_devices()?
            .find(|d| d.vendor_id() == RPI_VENDOR_ID && d.product_id() == DEBUG_PROBE_PRODUCT_ID)
            .ok_or_else(|| anyhow!("Could not find debug probe"))?;

        let device = device_info.open()?;
        let interface = device.claim_interface(0)?;

        debug!("{:?}", PigweedVendorCommand::Hello as u16);
        let result = block_on(interface.control_in(ControlIn {
            control_type: ControlType::Vendor,
            recipient: Recipient::Device,
            request: 0x50,
            value: 0x0,
            index: PigweedVendorCommand::Hello as u16,
            length: 256,
        }));
        if let Err(e) = result.status {
            return Err(anyhow!("Could not query Pigweed vendor interface: {e}"));
        }
        if &result.data != b"PIGWEED" {
            info!("{result:?}");

            return Err(anyhow!(
                "Pigweed vendor interface did not respond correctly"
            ));
        }

        Ok(Self { interface })
    }

    pub fn set_target(&mut self, target: SwdTarget) -> Result<()> {
        let value = match target {
            SwdTarget::Dut => 0,
            SwdTarget::Buddy => 1,
        };
        let result = block_on(self.interface.control_out(ControlOut {
            control_type: ControlType::Vendor,
            recipient: Recipient::Device,
            request: 0x50,
            value,
            index: PigweedVendorCommand::SetTarget as u16,
            data: &[],
        }));
        if let Err(e) = result.status {
            return Err(anyhow!("Failed to set target to {target:?}: {e}"));
        }

        Ok(())
    }

    pub fn get_info(&mut self) -> Result<TestTargetInfo> {
        let result = block_on(self.interface.control_in(ControlIn {
            control_type: ControlType::Vendor,
            recipient: Recipient::Device,
            request: 0x50,
            value: 0x0,
            index: PigweedVendorCommand::GetTargetInfo as u16,
            length: 256,
        }));
        if let Err(e) = result.status {
            return Err(anyhow!("Could not query Pigweed target info: {e}"));
        }
        let json = str::from_utf8(&result.data)?;
        debug!("json info string: {json}");
        let info = serde_json::from_str(json)?;

        Ok(info)
    }

    pub fn set_led(&mut self, red: u8, green: u8, blue: u8) -> Result<()> {
        let result = block_on(self.interface.control_out(ControlOut {
            control_type: ControlType::Vendor,
            recipient: Recipient::Device,
            request: 0x50,
            value: 0x0,
            index: PigweedVendorCommand::SetLed as u16,
            data: &[red, green, blue],
        }));

        if let Err(e) = result.status {
            return Err(anyhow!("Failed to set target to test status LED: {e}"));
        }

        Ok(())
    }
}

fn info_command() -> Result<()> {
    let mut target = PigweedTestTarget::new()?;
    let info = target.get_info()?;
    println!("Board Name: {}", info.name);
    println!("Microcontroller: {}", info.mcu);
    println!("Capabilities:");
    for cap in &info.capabilities {
        println!(" - {cap:?}");
    }

    Ok(())
}

fn set_led_command(red: u8, green: u8, blue: u8) -> Result<()> {
    let mut target = PigweedTestTarget::new()?;
    target.set_led(red, green, blue)?;
    println!("Target status LED set to #{red:02x}{green:02x}{blue:02x}");
    Ok(())
}

fn target_command(swd_target: SwdTarget) -> Result<()> {
    let mut target = PigweedTestTarget::new()?;
    target.set_target(swd_target)?;
    println!("SWD Target set to {swd_target:?}");
    Ok(())
}

fn main() {
    env_logger::init();
    let cli = Cli::parse();
    let result = match cli.command {
        Commands::Info => info_command(),
        Commands::SetLed { red, green, blue } => set_led_command(red, green, blue),
        Commands::Target { target } => target_command(target),
    };
    if let Err(e) = result {
        println!("command failed: {e}");
        std::process::exit(1);
    }
}