path: root/src/vtcol.rs

use anyhow::{anyhow, Result};
use std::{convert::TryFrom,
          fmt,
          io::Error,
          path::{Path, PathBuf},
          sync::atomic::{AtomicBool, Ordering}};

type Fd = libc::c_int;

static VERBOSITY: AtomicBool = AtomicBool::new(false);

macro_rules! vrb {
    ( $( $e:expr ),* ) => {(
            if VERBOSITY.load(Ordering::SeqCst) { println!( $( $e ),* ) }
    )}
}

const PALETTE_SIZE: usize = 16;
const PALETTE_BYTES: usize = PALETTE_SIZE * 3; /* 16 * sizeof(int) */

const RAW_COLEXPR_SIZE: usize = 6; /* e. g. 0xBADF00 */

type RawPalette<'a> = [&'a str; PALETTE_SIZE];

/* XXX: can we get these into ``libc``? */
const KDGKBTYPE: libc::c_int = 0x4b33; /* kd.h */
const GIO_CMAP: libc::c_int = 0x00004B70; /* kd.h */
const PIO_CMAP: libc::c_int = 0x00004B71; /* kd.h */
const KB_101: libc::c_char = 0x0002; /* kd.h */
const O_NOCTTY: libc::c_int = 0o0400; /* fcntl.h */

#[derive(Debug)]
enum Color
{
    Black(bool),
    Red(bool),
    Green(bool),
    Yellow(bool),
    Blue(bool),
    Magenta(bool),
    Cyan(bool),
    White(bool),
}

impl TryFrom<u8> for Color
{
    type Error = anyhow::Error;

    fn try_from(val: u8) -> Result<Self, Self::Error>
    {
        match val {
            0x00 => Ok(Color::Black(false)),
            0x01 => Ok(Color::Red(false)),
            0x02 => Ok(Color::Green(false)),
            0x03 => Ok(Color::Yellow(false)),
            0x04 => Ok(Color::Blue(false)),
            0x05 => Ok(Color::Magenta(false)),
            0x06 => Ok(Color::Cyan(false)),
            0x07 => Ok(Color::White(false)),

            0x08 => Ok(Color::Black(true)),
            0x09 => Ok(Color::Red(true)),
            0x0a => Ok(Color::Green(true)),
            0x0b => Ok(Color::Yellow(true)),
            0x0c => Ok(Color::Blue(true)),
            0x0d => Ok(Color::Magenta(true)),
            0x0e => Ok(Color::Cyan(true)),
            0x0f => Ok(Color::White(true)),

            _ => Err(anyhow!("invalid color value: {}", val)),
        }
    }
} /* [impl TryFrom<u8> for Color] */

impl Color
{
    fn format_brightness(b: bool, s: &str) -> String
    {
        if b {
            "bright ".to_string() + s
        } else {
            s.to_string()
        }
    }
} /* [impl Color] */

impl fmt::Display for Color
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result
    {
        let c = match *self {
            Color::Black(b) => Color::format_brightness(b, "black"),
            Color::Red(b) => Color::format_brightness(b, "red"),
            Color::Green(b) => Color::format_brightness(b, "green"),
            Color::Yellow(b) => Color::format_brightness(b, "yellow"),
            Color::Blue(b) => Color::format_brightness(b, "blue"),
            Color::Magenta(b) => Color::format_brightness(b, "magenta"),
            Color::Cyan(b) => Color::format_brightness(b, "cyan"),
            Color::White(b) => Color::format_brightness(b, "white"),
        };

        write!(f, "{}", c)
    }
} /* [impl fmt::Display for Color] */

#[derive(Debug, Clone)]
struct Builtin
{
    names:   &'static [&'static str],
    palette: &'static RawPalette<'static>,
}

const BUILTIN_SCHEMES: &[Builtin] = &[
    Builtin::solarized(),
    Builtin::solarized_light(),
    Builtin::default(),
    Builtin::phosphor(),
];

impl Builtin
{
    fn name(&self) -> &'static str { self.names.iter().next().unwrap() }

    const fn solarized() -> Self
    {
        Self {
            names:   &["solarized", "solarized_dark", "sd"],
            palette: &SOLARIZED_COLORS_DARK,
        }
    }

    const fn solarized_light() -> Self
    {
        Self {
            names:   &["solarized_light", "sl"],
            palette: &SOLARIZED_COLORS_LIGHT,
        }
    }

    const fn default() -> Self
    {
        Self {
            names:   &["default", "normal", "linux"],
            palette: &DEFAULT_COLORS,
        }
    }

    const fn phosphor() -> Self
    {
        Self { names: &["phosphor", "matrix"], palette: &MONOCHROME_PHOSPHOR }
    }
}

impl TryFrom<&str> for Builtin
{
    type Error = anyhow::Error;

    fn try_from(name: &str) -> Result<Self, Self::Error>
    {
        for b in BUILTIN_SCHEMES {
            if b.names.contains(&name) {
                return Ok(b.clone());
            }
        }

        Err(anyhow!("no such builtin: {}", name))
    }
}

impl<'a> fmt::Display for Builtin
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result
    {
        write!(f, "{}", self.name())
    }
}

#[derive(Debug)]
enum Scheme
{
    /** One of the predefined schemes. */
    Builtin(Builtin),
    /** Custom ``Palette``. */
    Palette(Palette),
    /** Load from file. */
    Custom(Option<PathBuf>),
}

impl<'a> fmt::Display for Scheme
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result
    {
        match self {
            Self::Builtin(b) => write!(f, "{}", b),
            Self::Custom(None) => write!(f, "<read stdin>"),
            Self::Custom(Some(fname)) => write!(f, "{}", fname.display()),
            Self::Palette(pal) => write!(f, "palette: {}", pal),
        }
    }
} /* [impl fmt::String for Scheme] */

impl Scheme
{
    fn from_stdin() -> Self { Self::Custom(None) }

    fn from_path<P: AsRef<Path>>(path: P) -> Self
    {
        Self::Custom(Some(path.as_ref().into()))
    }
} /* [impl Scheme] */

/** Try to select one of the predefined schemes; if that fails,
interpret the argument as a path. */
impl From<&str> for Scheme
{
    fn from(name: &str) -> Scheme
    {
        Builtin::try_from(name)
            .map(Self::Builtin)
            .unwrap_or_else(|_| Self::from_path(name))
    }
}

/** Try to match the palette against one of the predefined schemes. */
impl From<Palette> for Scheme
{
    fn from(pal: Palette) -> Scheme
    {
        if pal == Palette::from(&DEFAULT_COLORS) {
            return Self::Builtin(Builtin::default());
        }
        if pal == Palette::from(&SOLARIZED_COLORS_DARK) {
            return Self::Builtin(Builtin::solarized());
        }
        if pal == Palette::from(&SOLARIZED_COLORS_LIGHT) {
            return Self::Builtin(Builtin::solarized_light());
        }
        if pal == Palette::from(&MONOCHROME_PHOSPHOR) {
            return Self::Builtin(Builtin::phosphor());
        }

        Self::Palette(pal)
    }
}

/* struct Job -- Runtime parameters.
 */
#[derive(Debug)]
enum Job
{
    /** List available schemes. */
    List,
    /** Dump a scheme. */
    Dump(Scheme),
    /** Switch to color scheme. */
    Set(Scheme),
    /** Get currently active scheme. */
    Get,
    /** Toggle between two schemes. */
    Toggle(Scheme, Scheme),
}

impl<'a> Job
{
    pub fn from_argv() -> Result<Job>
    {
        use clap::{App, Arg, SubCommand};

        let app = App::new(clap::crate_name!())
            .version(clap::crate_version!())
            .author(clap::crate_authors!())
            .about(clap::crate_description!())
            .subcommand(
                SubCommand::with_name("dump").about("dump a color scheme").arg(
                    Arg::with_name("scheme")
                        .help("name of the scheme")
                        .required(true)
                        .value_name("NAME")
                        .takes_value(true)
                ),
            )
            .subcommand(
                SubCommand::with_name("list").about("list available schemes"),
            )
            .subcommand(
                SubCommand::with_name("set")
                    .about("apply color scheme to current terminal")
                    .arg(
                        Arg::with_name("scheme")
                            .value_name("NAME")
                            .help("predefined color scheme")
                            .takes_value(true)
                            .conflicts_with("file"),
                    )
                    .arg(
                        Arg::with_name("file")
                            .short("f")
                            .long("file")
                            .value_name("PATH")
                            .help("apply scheme from file")
                            .takes_value(true),
                    ),
            )
            .subcommand(
                SubCommand::with_name("get").about("get current color scheme"),
            )
            .subcommand(
                SubCommand::with_name("toggle")
                    .about("toggle between two schemes")
                    .arg(
                        Arg::with_name("one")
                            .value_name("NAME1")
                            .help("predefined color scheme")
                            .takes_value(true),
                    )
                    .arg(
                        Arg::with_name("two")
                            .value_name("NAME2")
                            .help("predefined color scheme")
                            .takes_value(true),
                    ),
            )
            .arg(
                Arg::with_name("verbose")
                    .short("v")
                    .long("verbose")
                    .help("enable extra diagnostics")
                    .takes_value(false),
            );

        let matches = app.get_matches();

        if matches.is_present("verbose") {
            VERBOSITY.store(true, Ordering::SeqCst);
        }

        match matches.subcommand() {
            ("dump", Some(subm)) => {
                if let Some(name) = subm.value_of("scheme") {
                    let scm = Scheme::from(name);
                    return Ok(Self::Dump(scm));
                }
                Err(anyhow!("dump requires an argument"))
            },
            ("list", _) => Ok(Self::List),
            ("set", Some(subm)) => {
                let scheme = match subm.value_of("scheme") {
                    Some("-") => Scheme::from_stdin(),
                    Some(name) => {
                        vrb!("pick predefined scheme [{}]", name);
                        Scheme::from(name)
                    },
                    None =>
                        match subm.value_of("file") {
                            None | Some("-") => Scheme::from_stdin(),
                            Some(fname) => {
                                vrb!(
                                    "read custom scheme from file [{}]",
                                    fname
                                );
                                Scheme::from_path(fname)
                            },
                        },
                };
                Ok(Self::Set(scheme))
            },
            ("get", _) => Ok(Self::Get),
            ("toggle", Some(subm)) => {
                match (subm.value_of("one"), subm.value_of("two")) {
                    (Some(one), Some(two)) => {
                        vrb!("toggle schemes [{}] and [{}]", one, two);
                        Ok(Self::Toggle(Scheme::from(one), Scheme::from(two)))
                    },
                    _ =>
                        Err(anyhow!(
                            "please supply two schemes to toggle between"
                        )),
                }
            },
            (junk, _) =>
                Err(anyhow!(
                    "invalid subcommand [{}]; try ``{} --help``",
                    junk,
                    clap::crate_name!()
                )),
        }
    }

    fn list_schemes()
    {
        println!("{} color schemes available:", BUILTIN_SCHEMES.len());
        for s in BUILTIN_SCHEMES {
            println!("      * {}", s.name());
        }
    }

    fn dump(scm: Scheme)
    {
        vrb!("Dumping color scheme {}", scm);
        match scm {
            Scheme::Builtin(Builtin { palette, .. }) =>
                Self::dump_scheme(palette),
            Scheme::Custom(None) => Self::dump_palette(Palette::from_stdin()),
            Scheme::Custom(Some(fname)) =>
                Self::dump_palette(Palette::from_file(&fname)),
            Scheme::Palette(pal) => Self::dump_palette(pal),
        }
    }

    fn dump_scheme(colors: &[&str; PALETTE_SIZE])
    {
        let pal: Palette = Palette::from(colors);
        pal.dump();
    }

    fn dump_palette(pal: Palette) { pal.dump() }

    fn run(self) -> Result<()>
    {
        match self {
            Self::Dump(scm) => Self::dump(scm),
            Self::List => Self::list_schemes(),
            Self::Set(scm) => Self::set_scheme(scm)?,
            Self::Get => Self::get_scheme()?,
            Self::Toggle(one, two) => Self::toggle_scheme(one, two)?,
        }

        Ok(())
    }

    fn set_scheme(scheme: Scheme) -> Result<()>
    {
        let pal = Palette::from(&scheme);
        vrb!("Using palette:");
        vrb!("{}", pal);
        let fd = get_console_fd()?;
        vrb!("console fd: {}", fd);

        ioctl_pio_cmap(fd, &pal)?;

        clear_term(fd)?;
        vrb!("successfully enabled scheme {:?}", scheme);
        /* It’s fine to leak the fd, the kernel will clean up anyways. */
        Ok(())
    }

    fn get_scheme() -> Result<()>
    {
        let fd = get_console_fd()?;
        vrb!("console fd: {}", fd);

        let pal = ioctl_gio_cmap(fd)?;
        let scm = Scheme::from(pal);

        vrb!("active scheme:");
        println!("{}", scm);

        Ok(())
    }

    /** Toggle between two schemes. Defaults to ``one`` in case neither scheme
    is active.
    */
    fn toggle_scheme(one: Scheme, two: Scheme) -> Result<()>
    {
        let fd = get_console_fd()?;
        vrb!("console fd: {}", fd);

        let pal = ioctl_gio_cmap(fd)?;

        if pal == Palette::from(&one) {
            Self::set_scheme(two)
        } else {
            Self::set_scheme(one)
        }
    }
} /* [impl Job] */

/* Rust appears to come with two wrappers for ``ioctl(2)``, but neither can be utilized for our
 * purposes. The one in ``sys`` is part of a private (seriously‽) whereas the one in the
 * ``libc`` module is defined as taking variable arguments and therefore cannot be called from
 * Rust. Wrapping C is still a bit awkward, as it seems.
 */
extern "C" {
    pub fn ioctl(
        d: libc::c_int,
        request: libc::c_int,
        data: *mut libc::c_void,
    ) -> libc::c_int;
}

const CONSOLE_PATHS: [&str; 6] = [
    "/proc/self/fd/0",
    "/dev/tty",
    "/dev/tty0",
    "/dev/vc/0",
    "/dev/systty",
    "/dev/console",
];

/** Vanilla Linux colors. */
const DEFAULT_COLORS: RawPalette = [
    "000000", "aa0000", "00aa00", "aa5500", "0000aa", "aa00aa", "00aaaa",
    "aaaaaa", "555555", "ff5555", "55ff55", "ffff55", "5555ff", "ff55ff",
    "55ffff", "ffffff",
];

/** The dark (default) version of the Solarized scheme. */
const SOLARIZED_COLORS_DARK: RawPalette = [
    "002b36", "dc322f", "859900", "b58900", "268bd2", "d33682", "2aa198",
    "eee8d5", "002b36", "cb4b16", "586e75", "657b83", "839496", "6c71c4",
    "93a1a1", "fdf6e3",
];

/** The light version of the Solarized theme. */
const SOLARIZED_COLORS_LIGHT: RawPalette = [
    "eee8d5", "dc322f", "859900", "b58900", "268bd2", "d33682", "2aa198",
    "073642", "fdf6e3", "cb4b16", "93a1a1", "839496", "657b83", "6c71c4",
    "586e75", "002b36",
];

/** Bright green monochrome terminal. */
const MONOCHROME_PHOSPHOR: RawPalette = [
    "000000", "68fc68", "68fc68", "68fc68", "68fc68", "68fc68", "68fc68",
    "68fc68", "68fc68", "68fc68", "68fc68", "68fc68", "68fc68", "68fc68",
    "68fc68", "68fc68",
];

const DUMMY_COLORS: RawPalette = [
    "000000", "ffffff", "000000", "ffffff", "000000", "ffffff", "000000",
    "ffffff", "000000", "ffffff", "000000", "ffffff", "000000", "ffffff",
    "000000", "ffffff",
];

fn nibble_of_char(chr: u8) -> u8
{
    match chr {
        b'0'..=b'9' => chr - b'0',
        b'a'..=b'f' => chr - b'a' + 10,
        b'A'..=b'F' => chr - b'A' + 10,
        _ => 0,
    }
}

macro_rules! byte_of_hex {
    ($ar:ident, $off:expr) => {
        (nibble_of_char($ar[$off])) << 4 | nibble_of_char($ar[$off + 1]) as u8
    };
}

fn rgb_of_hex_triplet(def: &str) -> (u8, u8, u8)
{
    let bytes = def.as_bytes();
    let r: u8 = byte_of_hex!(bytes, 0);
    let g: u8 = byte_of_hex!(bytes, 2);
    let b: u8 = byte_of_hex!(bytes, 4);
    (r, g, b)
}

#[derive(Eq, PartialEq, Clone)]
pub struct Palette([u8; PALETTE_BYTES]);

impl Palette
{
    /** Construct an all-zero ``Palette``. */
    pub fn new() -> Self { Self([0u8; PALETTE_BYTES]) }

    pub fn dummy() -> Self { Self::from(&DUMMY_COLORS) }

    pub fn from_buffered_reader(reader: &mut dyn std::io::BufRead) -> Self
    {
        let mut pal_idx: usize = 0;
        let mut pal: [u8; PALETTE_BYTES] = [0; PALETTE_BYTES];
        let mut line: String = String::new();

        while reader.read_line(&mut line).is_ok() {
            let len = line.len();
            if len == 0 {
                break;
            } else if len >= 8 {
                if let Some(off) = line.find('#') {
                    if off != 0 {
                        /* Palette index specified, number prepended */
                        let parse_res: Result<usize, _> =
                            std::str::FromStr::from_str(&line[0..off]);
                        if let Ok(new_idx) = parse_res {
                            if new_idx < PALETTE_SIZE {
                                pal_idx = new_idx * 3;
                            }
                        }
                    }
                    let off = off + 1;
                    if off > len - 6 {
                        /* no room left for color definition after '#' char */
                        panic!("invalid color definition: {}", line);
                    }
                    let col = &line[off..(off + RAW_COLEXPR_SIZE)];

                    let (r, g, b) = rgb_of_hex_triplet(col);
                    pal[pal_idx] = r;
                    pal[pal_idx + 1] = g;
                    pal[pal_idx + 2] = b;
                    pal_idx = (pal_idx + 3) % PALETTE_BYTES;
                }
            }
            line.truncate(0);
        }

        Self(pal)
    }

    fn dump(&self)
    {
        let mut buf: [u8; 3] = [0u8, 0u8, 0u8];
        for (i, col) in self.0.iter().enumerate() {
            let idx: usize = i % 3;
            buf[idx] = *col;
            if idx == 2 {
                let col = Color::try_from((i / 3) as u8)
                    .map(|c| format!("{}", c))
                    .unwrap_or_else(|_| "??".into());
                println!(
                    "{:>15} => 0x{:02.X}{:02.X}{:02.X}",
                    col, buf[0], buf[1], buf[2]
                );
            }
        }
    }

    pub fn from_file(fname: &Path) -> Self
    {
        /* Check if file exists
         */
        let file = match std::fs::File::open(&fname) {
            Err(e) => {
                panic!("failed to open {} as file ({})", fname.display(), e)
            },
            Ok(f) => f,
        };
        let mut reader = std::io::BufReader::new(file);

        /* Parse scheme file
         */
        Self::from_buffered_reader(&mut reader)
    }

    /* [Palette::from_file] */

    pub fn from_stdin() -> Self
    {
        vrb!("Go ahead, type your color scheme …");
        vrb!("vtcol>");
        let mut reader = std::io::BufReader::new(std::io::stdin());

        /* Parse scheme file
         */
        Self::from_buffered_reader(&mut reader)
    } /* [Palette::from_stdin] */
} /* [impl Palette] */

impl fmt::Display for Palette
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result
    {
        let mut i = 0;
        while i < PALETTE_BYTES {
            let _ = write!(f, "{}", if i == 0 { "(" } else { "\n " });
            let r = self.0[i];
            let g = self.0[i + 1];
            let b = self.0[i + 2];
            let _ = write!(
                f,
                "((r 0x{:02.X}) (g 0x{:02.X}) (b 0x{:02.x}))",
                r, g, b
            );
            i += 3;
        }
        writeln!(f, ")")
    }
} /* [impl fmt::Display for Palette] */

impl fmt::Debug for Palette
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result
    {
        let mut i: u8 = 0_u8;
        while (i as usize) < PALETTE_BYTES {
            let r = self.0[i as usize];
            let g = self.0[i as usize + 1];
            let b = self.0[i as usize + 2];
            let col = Color::try_from((i / 3) as u8)
                .map(|c| format!("{}", c))
                .unwrap_or_else(|_| "??".into());
            let _ = writeln!(f, "{} => 0x{:02.X}{:02.X}{:02.X}", col, r, g, b);
            i += 3_u8;
        }
        std::result::Result::Ok(())
    }
} /* [impl fmt::Debug for Palette] */

/** Obtain a ``Palette`` from a ``Scheme``. */
impl From<&Scheme> for Palette
{
    fn from(scm: &Scheme) -> Self
    {
        match scm {
            Scheme::Builtin(Builtin { palette, .. }) => Self::from(*palette),
            Scheme::Custom(None) => Self::from_stdin(),
            Scheme::Custom(Some(ref fname)) => Self::from_file(fname),
            Scheme::Palette(pal) => pal.clone(),
        }
    }
}

/** Obtain a ``Palette`` from a ``RawPalette``. */
impl From<&RawPalette<'_>> for Palette
{
    fn from(colors: &RawPalette<'_>) -> Self
    {
        let mut idx: usize = 0;
        let mut pal: [u8; PALETTE_BYTES] = [0; PALETTE_BYTES];

        for def in colors.iter() {
            let (r, g, b) = rgb_of_hex_triplet(*def);
            pal[idx] = r;
            pal[idx + 1] = g;
            pal[idx + 2] = b;
            //println!(">> {} -> {:X} {:X} {:X}", def, r, g, b);
            idx += 3;
        }

        Self(pal)
    }
}

fn fd_of_path(path: &std::path::Path) -> Option<Fd>
{
    let p = std::ffi::CString::new(path.to_str().unwrap()).unwrap();
    match unsafe { libc::open(p.as_ptr(), libc::O_RDWR | O_NOCTTY, 0) } {
        -1 => None,
        fd => {
            vrb!("  *> got fd");
            if unsafe { libc::isatty(fd) } == 0 {
                vrb!("  *> not a tty");
                return None;
            }

            let mut tty_type: libc::c_char = 0;

            let res = unsafe {
                ioctl(
                    fd,
                    KDGKBTYPE as libc::c_int,
                    std::mem::transmute(&mut tty_type),
                )
            };
            if res < 0 {
                vrb!("  *> ioctl failed");
                return None;
            }

            if tty_type != KB_101 {
                return None;
            }

            Some(fd)
        },
    }
}

fn get_console_fd() -> Result<Fd>
{
    for path in CONSOLE_PATHS.iter() {
        vrb!("trying path: {:?}", path);
        let path = std::path::Path::new(path);
        if let Some(fd) = fd_of_path(path) {
            vrb!(" * Success!");
            return Ok(fd);
        }
    }
    Err(anyhow!("could not retrieve fd for any of the search paths"))
}

fn write_to_term(fd: Fd, buf: &str) -> Result<()>
{
    let len = buf.len() as libc::size_t;

    if unsafe { libc::write(fd, buf.as_ptr() as *const libc::c_void, len) }
        != len as isize
    {
        Err(anyhow!(
            "failed to write {} B to fd {}: {}",
            len,
            fd,
            Error::last_os_error()
        ))
    } else {
        Ok(())
    }
}

fn clear_term(fd: Fd) -> Result<()>
{
    let clear = "\x1b[2J";
    let cursor = "\x1b[1;1H";
    write_to_term(fd, clear)?;
    write_to_term(fd, cursor)?;

    Ok(())
}

fn ioctl_pio_cmap(fd: Fd, pal: &Palette) -> Result<()>
{
    if unsafe { ioctl(fd, PIO_CMAP, std::mem::transmute(pal)) } < 0 {
        Err(anyhow!(
            "ioctl(PIO_CMAP) failed to insert new palette: {}",
            Error::last_os_error()
        ))
    } else {
        Ok(())
    }
}

fn ioctl_gio_cmap(fd: Fd) -> Result<Palette>
{
    let mut pal = Palette::new();

    if unsafe { ioctl(fd, GIO_CMAP, std::mem::transmute(&mut pal)) } < 0 {
        Err(anyhow!(
            "ioctl(GIO_CMAP) failed to get current palette: {}",
            Error::last_os_error()
        ))
    } else {
        Ok(pal)
    }
}

fn main() -> Result<()>
{
    let job = Job::from_argv()?;
    vrb!("job parms: {:?}", job);

    job.run()
}