ANSIColour.hs

   1 -- This file is part of Diohsc
   2 -- Copyright (C) 2020 Martin Bays <mbays@sdf.org>
   3 --
   4 -- This program is free software: you can redistribute it and/or modify
   5 -- it under the terms of version 3 of the GNU General Public License as
   6 -- published by the Free Software Foundation, or any later version.
   7 --
   8 -- You should have received a copy of the GNU General Public License
   9 -- along with this program.  If not, see http://www.gnu.org/licenses/.
  10
  11 {-# LANGUAGE OverloadedStrings #-}
  12 {-# LANGUAGE Safe              #-}
  13
  14 -- Basic ansi attributes, using only most widely supported ansi terminal codes
  15 module ANSIColour
  16     ( applyIf
  17     , resetCode
  18     , withColour
  19     , withBold
  20     , withReverse
  21     , withUnderline
  22     , withColourStr
  23     , withBoldStr
  24     , withReverseStr
  25     , withUnderlineStr
  26     , stripCSI
  27     , visibleLength
  28     , escapePromptCSI
  29     , sanitiseNonCSI
  30     , Colour(..)
  31     ) where
  32
  33 import           Control.Exception.Base (bracket_)
  34
  35 import qualified Data.Text.Lazy         as T
  36 import qualified Data.Text.Lazy.IO      as T
  37
  38 import           MetaString
  39
  40 data Colour = Black | Red | Green | Yellow
  41     | Blue | Magenta | Cyan | White
  42     | BoldBlack | BoldRed | BoldGreen | BoldYellow
  43     | BoldBlue | BoldMagenta | BoldCyan | BoldWhite
  44     deriving (Eq,Ord,Show,Read)
  45
  46 resetCode, boldCode, unboldCode, reverseCode,
  47     unreverseCode, underlineCode , ununderlineCode, resetColourCode
  48     :: MetaString a => a
  49 resetCode = "\ESC[0m"
  50 boldCode = "\ESC[1m"
  51 underlineCode = "\ESC[4m"
  52 reverseCode = "\ESC[7m"
  53 unboldCode = "\ESC[22m"
  54 ununderlineCode = "\ESC[24m"
  55 unreverseCode = "\ESC[27m"
  56 resetColourCode = "\ESC[39m\ESC[22m"
  57
  58 colourCode :: MetaString a => Colour -> a
  59 colourCode c = (if isBold c then boldCode else "") <> "\ESC[3" <> fromString (colNum c) <> "m"
  60     where
  61     isBold = flip elem [BoldBlack, BoldRed, BoldGreen, BoldYellow,
  62         BoldBlue, BoldMagenta, BoldCyan, BoldWhite]
  63     colNum Black       = "0"
  64     colNum Red         = "1"
  65     colNum Green       = "2"
  66     colNum Yellow      = "3"
  67     colNum Blue        = "4"
  68     colNum Magenta     = "5"
  69     colNum Cyan        = "6"
  70     colNum White       = "7"
  71     colNum BoldBlack   = "0"
  72     colNum BoldRed     = "1"
  73     colNum BoldGreen   = "2"
  74     colNum BoldYellow  = "3"
  75     colNum BoldBlue    = "4"
  76     colNum BoldMagenta = "5"
  77     colNum BoldCyan    = "6"
  78     colNum BoldWhite   = "7"
  79
  80 withStyle :: T.Text -> T.Text -> IO a -> IO a
  81 withStyle c r = T.putStr c `bracket_` T.putStr r
  82 withColour :: Colour -> IO a -> IO a
  83 withColour c = withStyle (colourCode c) resetColourCode
  84 withBold, withReverse, withUnderline :: IO a -> IO a
  85 withBold = withStyle boldCode unboldCode
  86 withReverse = withStyle reverseCode unreverseCode
  87 withUnderline = withStyle underlineCode ununderlineCode
  88
  89 withStyleStr :: MetaString a => a -> a -> a -> a
  90 withStyleStr c r s = c <> s <> r
  91 withColourStr :: MetaString a => Colour -> a -> a
  92 withColourStr c = withStyleStr (colourCode c) resetColourCode
  93 withBoldStr, withReverseStr, withUnderlineStr :: MetaString a => a -> a
  94 withBoldStr = withStyleStr boldCode unboldCode
  95 withReverseStr = withStyleStr reverseCode unreverseCode
  96 withUnderlineStr = withStyleStr underlineCode ununderlineCode
  97
  98 -- |"applyIf cond f" is shorthand for "if cond then f else id"
  99 applyIf :: Bool -> (a -> a) -> (a -> a)
 100 applyIf True  = id
 101 applyIf False = const id
 102
 103
 104 endCSI :: Char -> Bool
 105 endCSI c = '@' <= c && c <= '~'
 106
 107 -- |strip all CSI escape sequences
 108 stripCSI :: T.Text -> T.Text
 109 stripCSI s =
 110     let (pre,post) = T.breakOn "\ESC[" s
 111     in if T.null post then pre
 112         else (pre <>) . stripCSI . T.drop 1 .
 113             T.dropWhile (not . endCSI) $ T.drop 2 post
 114
 115 visibleLength :: (Integral i) => T.Text -> i
 116 visibleLength = fromIntegral . T.length . stripCSI
 117
 118 -- |sanitise non-CSI escape sequences by turning \ESC into \\ESC
 119 -- (buggy terminals make these sequences a potential security hole;
 120 -- see e.g. https://nvd.nist.gov/vuln/detail/CVE-2020-9366 )
 121 sanitiseNonCSI :: T.Text -> T.Text
 122 sanitiseNonCSI s =
 123     let (pre,post) = T.breakOn "\ESC" s
 124     in if T.null post then pre else
 125         let post' = T.drop 1 post
 126             isCSI = T.take 1 post' == "["
 127         in pre <> (if isCSI then "\ESC" else "\\ESC") <> sanitiseNonCSI post'
 128
 129 -- |append \STX to each CSI sequence, as required in Haskeline prompts.
 130 -- See https://github.com/judah/haskeline/wiki/ControlSequencesInPrompt
 131 escapePromptCSI :: String -> String
 132 escapePromptCSI s = case break (== '\ESC') s of
 133     (pre,'\ESC':'[':post) -> ((pre <> "\ESC[") <>) $
 134         case break endCSI post of
 135             (pre',e:post') -> (pre' <>) $ e : '\STX' : escapePromptCSI post'
 136             (pre',[])      -> pre'
 137     (pre,[]) -> pre
 138     (pre,e:post) -> (pre <>) $ e : escapePromptCSI post