{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Tokstyle.Analysis.SecurityRank.Annotations
    ( AnnotationMap
    , FunctionAnnotation
    , parseAllAnnotations
    ) where

import           Control.Monad                          (when)
import           Control.Monad.State.Strict             (State, execState, get,
                                                         gets, modify, modify')
import           Data.Fix                               (Fix (..))
import           Data.Map.Strict                        (Map)
import qualified Data.Map.Strict                        as Map
import           Data.Maybe                             (isJust, mapMaybe)
import           Data.Text                              (Text)
import qualified Data.Text                              as T
import           Language.Cimple                        (CommentF (..),
                                                         NodeF (..))
import qualified Language.Cimple                        as C
import           Language.Cimple.TraverseAst            (AstActions (..),
                                                         astActions,
                                                         traverseAst)
import           Tokstyle.Analysis.SecurityRank.Lattice (SecurityRank (..))

-- | Maps a parameter name (or "return", "sink", "sink_max") to its rank.
type FunctionAnnotation = Map Text SecurityRank

-- | Maps a function or struct member name to its full annotation.
type AnnotationMap = Map Text FunctionAnnotation

data CollectorState = CollectorState
    { CollectorState -> AnnotationMap
csAnnotations   :: AnnotationMap
    , CollectorState -> Maybe Text
csCurrentStruct :: Maybe Text -- Still needed to construct the full member name.
    }

-- | Extracts security rank annotations from a list of comment items.
extractRanksFromItems :: [C.Comment (C.Lexeme Text)] -> [(Text, SecurityRank)]
extractRanksFromItems :: [Comment (Lexeme Text)] -> [(Text, SecurityRank)]
extractRanksFromItems = (Comment (Lexeme Text) -> Maybe (Text, SecurityRank))
-> [Comment (Lexeme Text)] -> [(Text, SecurityRank)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((Comment (Lexeme Text) -> Maybe (Text, SecurityRank))
 -> [Comment (Lexeme Text)] -> [(Text, SecurityRank)])
-> (Comment (Lexeme Text) -> Maybe (Text, SecurityRank))
-> [Comment (Lexeme Text)]
-> [(Text, SecurityRank)]
forall a b. (a -> b) -> a -> b
$ \(Fix CommentF (Lexeme Text) (Comment (Lexeme Text))
item) -> case CommentF (Lexeme Text) (Comment (Lexeme Text))
item of
    DocSecurityRank (C.L AlexPosn
_ LexemeClass
_ Text
kw) Maybe (Lexeme Text)
mparam (C.L AlexPosn
_ LexemeClass
_ Text
rankStr) ->
        let
            rankVal :: Int
rankVal = String -> Int
forall a. Read a => String -> a
read (Text -> String
T.unpack Text
rankStr)
            key :: Text
key = case Maybe (Lexeme Text)
mparam of
                    Maybe (Lexeme Text)
Nothing              -> Text
kw
                    Just (C.L AlexPosn
_ LexemeClass
_ Text
param) -> Text
kw Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
":" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
param
        in
            (Text, SecurityRank) -> Maybe (Text, SecurityRank)
forall a. a -> Maybe a
Just (Text
key, Int -> SecurityRank
Rank Int
rankVal)
    CommentF (Lexeme Text) (Comment (Lexeme Text))
_ -> Maybe (Text, SecurityRank)
forall a. Maybe a
Nothing

-- | Safely extracts the list of items from a DocLine, returning empty for other types.
getDocLineItems :: C.Comment (C.Lexeme Text) -> [C.Comment (C.Lexeme Text)]
getDocLineItems :: Comment (Lexeme Text) -> [Comment (Lexeme Text)]
getDocLineItems (Fix (DocLine [Comment (Lexeme Text)]
items)) = [Comment (Lexeme Text)]
items
getDocLineItems Comment (Lexeme Text)
_                     = []

-- | Traverses the AST to find commented declarations and parse their annotations.
annotationActions :: AstActions (State CollectorState) Text
annotationActions :: AstActions (State CollectorState) Text
annotationActions = AstActions (State CollectorState) Text
forall (f :: * -> *) text. Applicative f => AstActions f text
astActions
    { doNode :: String
-> Node (Lexeme Text)
-> State CollectorState ()
-> State CollectorState ()
doNode = \String
_file Node (Lexeme Text)
node State CollectorState ()
act -> do
        Maybe Text
originalStruct <- (CollectorState -> Maybe Text) -> State CollectorState (Maybe Text)
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CollectorState -> Maybe Text
csCurrentStruct

        let mNewStruct :: Maybe Text
mNewStruct = case Node (Lexeme Text) -> NodeF (Lexeme Text) (Node (Lexeme Text))
forall (f :: * -> *). Fix f -> f (Fix f)
unFix Node (Lexeme Text)
node of
                C.Struct (C.L AlexPosn
_ LexemeClass
_ Text
name) [Node (Lexeme Text)]
_                     -> Text -> Maybe Text
forall a. a -> Maybe a
Just Text
name
                C.Typedef (Fix (C.Struct (C.L AlexPosn
_ LexemeClass
_ Text
name) [Node (Lexeme Text)]
_)) Lexeme Text
_ -> Text -> Maybe Text
forall a. a -> Maybe a
Just Text
name
                NodeF (Lexeme Text) (Node (Lexeme Text))
_                                             -> Maybe Text
forall a. Maybe a
Nothing

        -- If this node defines a new struct context, update the state before processing it and its children.
        Bool -> State CollectorState () -> State CollectorState ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Maybe Text -> Bool
forall a. Maybe a -> Bool
isJust Maybe Text
mNewStruct) (State CollectorState () -> State CollectorState ())
-> State CollectorState () -> State CollectorState ()
forall a b. (a -> b) -> a -> b
$ do
            (CollectorState -> CollectorState) -> State CollectorState ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CollectorState -> CollectorState) -> State CollectorState ())
-> (CollectorState -> CollectorState) -> State CollectorState ()
forall a b. (a -> b) -> a -> b
$ \CollectorState
st -> CollectorState
st { csCurrentStruct :: Maybe Text
csCurrentStruct = Maybe Text
mNewStruct }

        -- Process the current node using the potentially updated context.
        Node (Lexeme Text) -> State CollectorState ()
processNode Node (Lexeme Text)
node

        -- Traverse into children.
        State CollectorState ()
act

        -- After traversing, restore the original context.
        (CollectorState -> CollectorState) -> State CollectorState ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CollectorState -> CollectorState) -> State CollectorState ())
-> (CollectorState -> CollectorState) -> State CollectorState ()
forall a b. (a -> b) -> a -> b
$ \CollectorState
st -> CollectorState
st { csCurrentStruct :: Maybe Text
csCurrentStruct = Maybe Text
originalStruct }
    }

-- | Processes a single node to find and store annotations.
processNode :: C.Node (C.Lexeme Text) -> State CollectorState ()
processNode :: Node (Lexeme Text) -> State CollectorState ()
processNode (Fix (C.Commented (Fix (C.CommentInfo (Fix (C.DocComment [Comment (Lexeme Text)]
docLines)))) Node (Lexeme Text)
decl)) = do
    CollectorState
st <- StateT CollectorState Identity CollectorState
forall s (m :: * -> *). MonadState s m => m s
get
    let allItems :: [Comment (Lexeme Text)]
allItems = (Comment (Lexeme Text) -> [Comment (Lexeme Text)])
-> [Comment (Lexeme Text)] -> [Comment (Lexeme Text)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Comment (Lexeme Text) -> [Comment (Lexeme Text)]
getDocLineItems [Comment (Lexeme Text)]
docLines
    let ranks :: [(Text, SecurityRank)]
ranks = [Comment (Lexeme Text)] -> [(Text, SecurityRank)]
extractRanksFromItems [Comment (Lexeme Text)]
allItems
    let annotation :: Map Text SecurityRank
annotation = [(Text, SecurityRank)] -> Map Text SecurityRank
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Text, SecurityRank)]
ranks

    -- Handle function annotations
    case Node (Lexeme Text) -> NodeF (Lexeme Text) (Node (Lexeme Text))
forall (f :: * -> *). Fix f -> f (Fix f)
unFix Node (Lexeme Text)
decl of
        C.FunctionDefn Scope
_ (Fix (C.FunctionPrototype Node (Lexeme Text)
_ (C.L AlexPosn
_ LexemeClass
_ Text
name) [Node (Lexeme Text)]
_)) Node (Lexeme Text)
_ ->
            Text -> Map Text SecurityRank -> State CollectorState ()
addAnnotations Text
name Map Text SecurityRank
annotation
        C.FunctionDecl Scope
_ (Fix (C.FunctionPrototype Node (Lexeme Text)
_ (C.L AlexPosn
_ LexemeClass
_ Text
name) [Node (Lexeme Text)]
_)) ->
            Text -> Map Text SecurityRank -> State CollectorState ()
addAnnotations Text
name Map Text SecurityRank
annotation
        C.TypedefFunction (Fix (C.FunctionPrototype Node (Lexeme Text)
_ (C.L AlexPosn
_ LexemeClass
_ Text
name) [Node (Lexeme Text)]
_)) ->
            Text -> Map Text SecurityRank -> State CollectorState ()
addAnnotations Text
name Map Text SecurityRank
annotation
        NodeF (Lexeme Text) (Node (Lexeme Text))
_ -> () -> State CollectorState ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

    -- Handle struct member annotations
    case (Node (Lexeme Text) -> NodeF (Lexeme Text) (Node (Lexeme Text))
forall (f :: * -> *). Fix f -> f (Fix f)
unFix Node (Lexeme Text)
decl, CollectorState -> Maybe Text
csCurrentStruct CollectorState
st) of
        (C.MemberDecl (Fix (C.VarDecl Node (Lexeme Text)
_ (C.L AlexPosn
_ LexemeClass
_ Text
memberName) [Node (Lexeme Text)]
_)) Maybe (Lexeme Text)
_, Just Text
structName) -> do
            let fullMemberName :: Text
fullMemberName = Text
structName Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"." Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
memberName
            Text -> Map Text SecurityRank -> State CollectorState ()
addAnnotations Text
fullMemberName Map Text SecurityRank
annotation
        (NodeF (Lexeme Text) (Node (Lexeme Text)), Maybe Text)
_ -> () -> State CollectorState ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
processNode Node (Lexeme Text)
_ = () -> State CollectorState ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()


-- | Adds a list of parsed annotations for a given entity (function or struct member).
addAnnotations :: Text -> FunctionAnnotation -> State CollectorState ()
addAnnotations :: Text -> Map Text SecurityRank -> State CollectorState ()
addAnnotations Text
name Map Text SecurityRank
annotation =
    (CollectorState -> CollectorState) -> State CollectorState ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((CollectorState -> CollectorState) -> State CollectorState ())
-> (CollectorState -> CollectorState) -> State CollectorState ()
forall a b. (a -> b) -> a -> b
$ \CollectorState
st -> CollectorState
st { csAnnotations :: AnnotationMap
csAnnotations = Text -> Map Text SecurityRank -> AnnotationMap -> AnnotationMap
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Text
name Map Text SecurityRank
annotation (CollectorState -> AnnotationMap
csAnnotations CollectorState
st) }

-- | The main entry point. Traverses all translation units and returns the complete
-- annotation map.
parseAllAnnotations :: [(FilePath, [C.Node (C.Lexeme Text)])] -> AnnotationMap
parseAllAnnotations :: [(String, [Node (Lexeme Text)])] -> AnnotationMap
parseAllAnnotations [(String, [Node (Lexeme Text)])]
tus =
    let
        initialState :: CollectorState
initialState = AnnotationMap -> Maybe Text -> CollectorState
CollectorState AnnotationMap
forall k a. Map k a
Map.empty Maybe Text
forall a. Maybe a
Nothing
        finalState :: CollectorState
finalState = State CollectorState () -> CollectorState -> CollectorState
forall s a. State s a -> s -> s
execState (AstActions (State CollectorState) Text
-> [(String, [Node (Lexeme Text)])] -> State CollectorState ()
forall text a (f :: * -> *).
(TraverseAst text a, Applicative f) =>
AstActions f text -> a -> f ()
traverseAst AstActions (State CollectorState) Text
annotationActions [(String, [Node (Lexeme Text)])]
tus) CollectorState
initialState
    in
        CollectorState -> AnnotationMap
csAnnotations CollectorState
finalState