diohsc-0.1.3: ANSIColour.hs

-- This file is part of Diohsc
-- Copyright (C) 2020 Martin Bays <mbays@sdf.org>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of version 3 of the GNU General Public License as
-- published by the Free Software Foundation, or any later version.
--
-- You should have received a copy of the GNU General Public License
-- along with this program.  If not, see http://www.gnu.org/licenses/.

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Safe              #-}

-- Basic ansi attributes, using only most widely supported ansi terminal codes
module ANSIColour
    ( applyIf
    , resetCode
    , withColour
    , withBold
    , withUnderline
    , withColourStr
    , withBoldStr
    , withUnderlineStr
    , stripCSI
    , visibleLength
    , escapePromptCSI
    , sanitiseNonCSI
    , Colour(..)
    ) where

import           Control.Exception.Base (bracket_)

import qualified Data.Text.Lazy         as T
import qualified Data.Text.Lazy.IO      as T

import           MetaString

data Colour = Black | Red | Green | Yellow
    | Blue | Magenta | Cyan | White
    | BoldBlack | BoldRed | BoldGreen | BoldYellow
    | BoldBlue | BoldMagenta | BoldCyan | BoldWhite
    deriving (Eq,Ord,Show,Read)

resetCode,boldCode,underlineCode :: MetaString a => a
resetCode = "\ESC[0m"
boldCode = "\ESC[1m"
underlineCode = "\ESC[4m"
colourCode :: MetaString a => Colour -> a
colourCode c = (if isBold c then boldCode else "") <> "\ESC[3" <> fromString (colNum c) <> "m"
    where
    isBold = flip elem [BoldBlack, BoldRed, BoldGreen, BoldYellow,
        BoldBlue, BoldMagenta, BoldCyan, BoldWhite]
    colNum Black       = "0"
    colNum Red         = "1"
    colNum Green       = "2"
    colNum Yellow      = "3"
    colNum Blue        = "4"
    colNum Magenta     = "5"
    colNum Cyan        = "6"
    colNum White       = "7"
    colNum BoldBlack   = "0"
    colNum BoldRed     = "1"
    colNum BoldGreen   = "2"
    colNum BoldYellow  = "3"
    colNum BoldBlue    = "4"
    colNum BoldMagenta = "5"
    colNum BoldCyan    = "6"
    colNum BoldWhite   = "7"

thenReset :: IO () -> IO a -> IO a
thenReset = (`bracket_` T.putStr resetCode)

withColour :: Colour -> IO a -> IO a
withColour c = thenReset $ T.putStr (colourCode c)
withBold :: IO a -> IO a
withBold = thenReset $ T.putStr boldCode
withUnderline :: IO a -> IO a
withUnderline = thenReset $ T.putStr underlineCode

withColourStr :: MetaString a => Colour -> a -> a
withColourStr c s = colourCode c <> s <> resetCode
withBoldStr :: MetaString a => a -> a
withBoldStr s = boldCode <> s <> resetCode
withUnderlineStr :: MetaString a => a -> a
withUnderlineStr s = underlineCode <> s <> resetCode

-- |"applyIf cond f" is shorthand for "if cond then f else id"
applyIf :: Bool -> (a -> a) -> (a -> a)
applyIf True  = id
applyIf False = const id


endCSI :: Char -> Bool
endCSI c = '@' <= c && c <= '~'

-- |strip all CSI escape sequences
stripCSI :: T.Text -> T.Text
stripCSI s =
    let (pre,post) = T.breakOn "\ESC[" s
    in if T.null post then pre
        else (pre <>) . stripCSI . T.drop 1 .
            T.dropWhile (not . endCSI) $ T.drop 2 post

visibleLength :: (Integral i) => T.Text -> i
visibleLength = fromIntegral . T.length . stripCSI

-- |sanitise non-CSI escape sequences by turning \ESC into \\ESC
-- (buggy terminals make these sequences a potential security hole;
-- see e.g. https://nvd.nist.gov/vuln/detail/CVE-2020-9366 )
sanitiseNonCSI :: T.Text -> T.Text
sanitiseNonCSI s =
    let (pre,post) = T.breakOn "\ESC" s
    in if T.null post then pre else
        let post' = T.drop 1 post
            isCSI = T.take 1 post' == "["
        in pre <> (if isCSI then "\ESC" else "\\ESC") <> sanitiseNonCSI post'

-- |append \STX to each CSI sequence, as required in Haskeline prompts.
-- See https://github.com/judah/haskeline/wiki/ControlSequencesInPrompt
escapePromptCSI :: String -> String
escapePromptCSI s = case break (== '\ESC') s of
    (pre,'\ESC':'[':post) -> ((pre <> "\ESC[") <>) $
        case break endCSI post of
            (pre',e:post') -> (pre' <>) $ e : '\STX' : escapePromptCSI post'
            (pre',[])      -> pre'
    (pre,[]) -> pre
    (pre,e:post) -> (pre <>) $ e : escapePromptCSI post