module Main where

import Debug.Trace
import Codec.Archive.Zip
import qualified Data.ByteString.Lazy as B
import Control.Monad
import Data.List
import Data.Word
import Data.Bits
import Data.Array.Unboxed
import System.Console.GetOpt
import System.Environment
import Data.Maybe
import Text.Regex.PCRE
import Data.Time.Clock.POSIX

import Data.Binary
import Data.Binary.Get
import Data.Char

bToString   = map (chr . fromIntegral) . B.unpack
bFromString = B.pack   . map (fromIntegral . ord) 
onString f  = bFromString . f . bToString

data Access = Public | Private | Protected | Package deriving (Eq)

instance Show Access where
    show Public = "public"
    show Private = "private"
    show Protected = "protected"
    show Package = ""


data Class = Class {
    clsAccess :: Access,
    clsMembers :: [Member],
    clsName :: B.ByteString,
    clsIsInterface :: Bool
} deriving (Show)

data Member = Method { mAccess :: Access, mName :: B.ByteString, mSig :: B.ByteString }
            | Field  { mAccess :: Access, mName :: B.ByteString, mSig :: B.ByteString }
            deriving (Show)

-- Types related to cmdline arguments
 
data ClassFileSource = ClassFile { path :: FilePath    }
                     | JarFile   { path :: FilePath    }
                     | ClassPath { paths :: [ClassFileSource] }

data Location = InFile FilePath
              | InJar  { jarPath :: FilePath, innerPath :: FilePath }
              deriving (Show)

newtype SearchSource = SearchSource (Class->Bool)
data SearchTarget = SearchClass
                  | SearchMember (Member->Bool)

data Args = Args {
    dataSource   :: [ClassFileSource],
    searchSource :: SearchSource,
    searchTarget :: SearchTarget
}

data Result = FoundClass Location Class
            | FoundMember Location Class Member
            deriving (Show)


----------------------------------------------------------------
--                    ARGUMENTS PARSING                       --
----------------------------------------------------------------


data PlainArgs = PlainArgs {
    typeRegex :: Maybe String,
    memberRegex :: Maybe String,
    typeAccess :: Maybe String,
    memberAccess :: Maybe String,
    target :: Maybe String 
}
emptyArgs = PlainArgs Nothing Nothing Nothing Nothing Nothing

data MemberKind = FieldKind | MethodKind deriving (Eq)

parseArgs :: PlainArgs -> [String] -> Args
parseArgs a paths = Args dataSource (SearchSource typeFilter) searchTarget
    where searchTarget = case (memberRegex a,memberAccess a,target a) of 
                            (Nothing,Nothing,Nothing) -> SearchClass
                            _                         -> SearchMember memberFilter

          searchMembers = isJust (memberRegex a) 

          typeFilter = accessAndRegex (access (typeAccess a)) (typeRegex a) clsAccess clsName
          
          memberFilter mem = (memKind `maybeEq` kind mem)
                           && accessAndRegex (access (memberAccess a)) (memberRegex a) 
                                             mAccess mName mem
              where kind (Field  _ _ _) = FieldKind
                    kind (Method _ _ _) = MethodKind

          memKind = case (target a) of 
                        Just ('f':_) -> Just FieldKind 
                        Just ('m':_) -> Just MethodKind 
                        _            -> Nothing

          access a = case a of
                        Just "public"    -> Just Public
                        Just "private"   -> Just Private
                        Just "protected" -> Just Protected
                        Just "package"   -> Just Package
                        _                -> Nothing

          dataSource = map toDataSource paths
          toDataSource p | ".jar"   `isSuffixOf` p = JarFile p
                         | ".class" `isSuffixOf` p = ClassFile p
                         | ":"      `isInfixOf`  p = ClassPath . map toDataSource $ (==':') `unjoin` p

          accessAndRegex acc rx getAcc getName x =  (acc `maybeEq` (getAcc x)) 
                                                 && (nameMatchesRegex x)
              where nameMatchesRegex x = case rx of 
                                            Nothing -> True
                                            Just rx -> getName x =~ rx
          
          maybeEq Nothing  _ = True
          maybeEq (Just a) b = a==b
          
          unjoin :: (a->Bool) -> [a] -> [[a]]
          unjoin p s = go [] [] s
              where go res cur []     = reverse (cur:res)
                    go res cur (x:xs) | p x       = go (cur:res) [] xs
                                      | otherwise = go res (x:cur) xs

----------------------------------------------------------------
--                            MAIN                            --
----------------------------------------------------------------

main :: IO ()
main = timed $ do
    (opts, paths, errs) <- getOpt Permute options `liftM` getArgs
    let compose = foldr (.) id
    case errs of
        [] -> go (parseArgs (compose opts emptyArgs) paths)
        _  -> mapM_ putStrLn errs >> showHelp

timed :: IO a -> IO a
timed a = do
    t1 <- getPOSIXTime
    ra <- a
    t2 <- getPOSIXTime
    putStrLn . show $ diffTimes t1 t2
    return ra

diffTimes ta tb  =  i tb - i ta
i x = realToFrac x

    

options :: [OptDescr (PlainArgs -> PlainArgs)]
options = [ Option [&#039;c&#039;] []         (OptArg (s a -> a { typeRegex    = s }) "RX")   
                "Search for/in types whose name (with package) contains a match of this regex",
            Option [&#039;m&#039;] []         (OptArg (s a -> a { memberRegex  = s }) "RX")  
                "Search for members whose name contains a match of this regex",
            Option []    ["ca"]     (OptArg (s a -> a { typeAccess   = s })"ACCESS")  
                "Search for/in types having the specified access (public/private/protected/package)",
            Option []    ["ma"]     (OptArg (s a -> a { memberAccess = s }) "ACCESS") 
                "Search for members having the specified access (public/private/protected/package)",
            Option [&#039;t&#039;] ["target"] (OptArg (s a -> a { target       = s }) "TYPE")   
                "Search for members of the given type (field=f/method=m)"]

showHelp = putStrLn $ usageInfo usage options
    where usage="jarf [OPTION]... FILE... - Search for classes/methods/interfaces in JAR file(s)"

go :: Args -> IO ()
go (Args dataS searchS searchT) = do
    classes <- parseDataSource dataS
    mapM_ (putStrLn . present) . concatMap (search searchS searchT) $ classes

-- Results presentation

present :: Result -> String
present (FoundClass loc cls) = presentLoc loc ++ ": " ++ presentClass cls
present (FoundMember loc cls mem) = presentLoc loc ++ ": " ++ 
                                    presentClass cls ++ ": " ++ 
                                    show (mAccess mem) ++ " " ++ 
                                    bToString (mName mem) ++ " :: " ++ 
                                    bToString (mSig mem)

presentLoc (InFile path) = path
presentLoc (InJar jp ip) = jp ++ "!" ++ ip

presentClass (Class acc _ name False) = show acc ++ " class "     ++ bToString name
presentClass (Class acc _ name True)  = show acc ++ " interface " ++ bToString name

-- Search

search :: SearchSource -> SearchTarget -> (Location,Class) -> [Result]
search (SearchSource classP) SearchClass (loc,c) = 
    if (classP c) then [FoundClass loc c] else []
search (SearchSource classP) (SearchMember memP) (loc,c) =
    if (classP c) then [FoundMember loc c mem | mem <- clsMembers c, memP mem]
                  else []


-- Getting the actuall classes to search in

parseDataSource :: [ClassFileSource] -> IO [(Location, Class)]
parseDataSource = fmap concat . sequence . map parseDataSource&#039;
parseDataSource&#039; :: ClassFileSource -> IO [(Location,Class)]
parseDataSource&#039; (ClassFile path) = do
    f <- parseClassFile `liftM` B.readFile path
    return [(InFile path, f)]
parseDataSource&#039; (JarFile jar)    = do
    archive <- toArchive `liftM` B.readFile jar
    return [(InJar jar fileName, parseClassFile $ fromEntry entry)
            | entry <- zEntries archive,
              let fileName = eRelativePath entry,
              ".class" `isSuffixOf` fileName]


----------------------------------------------------------------
--                      CLASS FILE PARSING                    --
----------------------------------------------------------------

-- Constant pool stuff

data ConstantPool = ConstantPool { 
    cpPos    :: Array Word16 Word16,
    cpData   :: B.ByteString
}

getUTF8FromCP :: ConstantPool -> Word16 -> B.ByteString
getUTF8FromCP cp@(ConstantPool cpPos cpData) i = res
    where pos = cpPos!i
          begin = fromIntegral $ pos+3
          len   = fromIntegral $ getWord16FromCP cp (pos+1)
          res = B.take len (B.drop begin cpData)

getWord16FromCP (ConstantPool cpPos cpData) iData = 
    (do skip (fromIntegral iData); readWord16) `runGet` cpData

-- The very class file parser

readByte = get :: Get Word8
readWord16 = get :: Get Word16
readInt16 = fromIntegral `liftM` readWord16
readWord32 = get :: Get Word32

parseClassFile = runGet classFileParser

classFileParser = do
    skip (4+2+2) -- magic, minor_version, major_version
    cpSize      <- readWord16
    cp          <- parseConstantPool cpSize
    accessFlags <- readWord16
    thisClass   <- readWord16
    let classNameIndex = getWord16FromCP cp (1 + (cpPos cp)!thisClass)
    let className      = getUTF8FromCP   cp classNameIndex
    let access         = flagsToAccess accessFlags
    let isInterface    = (0 /= accessFlags .&. 0x0200)
    skip 2                    -- super_class
    ifCount     <- readInt16  -- interfaces_count
    skip (2*ifCount)          -- u2 interfaces[interfaces_count]
    fieldCount  <- readInt16
    fields      <- replicateM fieldCount (parseMember cp Field decryptType) 
    methodCount <- readInt16
    methods     <- replicateM methodCount (parseMember cp Method decryptSig)
    return (Class access (fields++methods) className isInterface)

parseMember cp ctor decryptor = do
    accessFlags <- readWord16
    nameIndex   <- readWord16
    descrIndex  <- readWord16
    attCount    <- readInt16
    replicateM attCount $ do readWord16; readWord32 >>= (skip.fromIntegral)
    return $ ctor (flagsToAccess accessFlags)
                  (getUTF8FromCP cp nameIndex)
                  (decryptor (getUTF8FromCP cp descrIndex))

parseConstantPool :: Word16 -> Get ConstantPool    
parseConstantPool n = do
    bs         <- getRemainingLazyByteString
    cpPosList  <- parseConstantPool&#039; (n-1) 0 []
    cpPosArray <- return $ array (1, genericLength cpPosList) $ zip [1..] cpPosList
    return (ConstantPool cpPosArray bs)


parseConstantPool&#039; :: Word16 -> Word16 -> [Word16] -> Get [Word16]
parseConstantPool&#039; 0 pos res = return $ reverse (0:res)
parseConstantPool&#039; n pos res = do
  tag <- readByte
  let skipRet n = do skip n; return n
  delta <- fromIntegral `liftM` case tag of
    -- Yes, tag 2 is unused.
    1  -> do len <- readInt16  -- CONSTANT_Utf8
             skip len
             return (len+2)   
    3  -> skipRet 4            -- CONSTANT_Integer
    4  -> skipRet 4            -- CONSTANT_Float
    5  -> skipRet 8            -- CONSTANT_Long (2 slots)
    6  -> skipRet 8            -- CONSTANT_Double (2 slots)
    7  -> skipRet 2            -- CONSTANT_Class
    8  -> skipRet 2            -- CONSTANT_String
    9  -> skipRet 4            -- CONSTANT_Fieldref
    10 -> skipRet 4            -- CONSTANT_Methodref
    11 -> skipRet 4            -- CONSTANT_InterfaceMethodref
    12 -> skipRet 4            -- CONSTANT_NameAndType
  let n&#039; = if (tag==5 || tag==6) then (n-2) else (n-1)
  case tag of
    5 -> parseConstantPool&#039; n&#039; (pos+delta+1) (pos:0:res)
    6 -> parseConstantPool&#039; n&#039; (pos+delta+1) (pos:0:res)
    _ -> parseConstantPool&#039; n&#039; (pos+delta+1) (pos:res)


-- Field/method type decryption

decryptType = onString (fst . decryptType&#039;)
decryptType&#039; ""      = ("","")
decryptType&#039; (a:s) = case a of
    &#039;B&#039; -> ("byte",s)
    &#039;C&#039; -> ("char",s)
    &#039;D&#039; -> ("double",s)
    &#039;F&#039; -> ("float",s)
    &#039;I&#039; -> ("int",s)
    &#039;J&#039; -> ("long",s)
    &#039;S&#039; -> ("short",s)
    &#039;Z&#039; -> ("boolean",s)
    &#039;V&#039; -> ("void",s)
    &#039;[&#039; -> let (t,s&#039;) = decryptType&#039; s in (t++"[]",s&#039;)
    &#039;L&#039; -> go "" s
        where go c (&#039;;&#039;:s) = (reverse c, s)
              go c (&#039;/&#039;:s) = go (&#039;.&#039;:c) s
              go c (a:s)   = go (a:c)   s

decryptSig = onString decryptSig&#039;
decryptSig&#039; s = concat (if null parTypes 
                        then ["()"]
                        else (intersperse " -> " parTypes)) 
                ++ " -> " ++ retType
    where (&#039;(&#039;:params, &#039;)&#039;:ret) = break (==&#039;)&#039;) s
          retType = fst (decryptType&#039; ret)
          parTypes = go [] params
          go ps "" = reverse ps
          go ps s  = let (p, s&#039;) = decryptType&#039; s in go (p:ps) s&#039;


flagsToAccess :: Word16 -> Access
flagsToAccess w | 0 /= w.&.0x0001 = Public
                | 0 /= w.&.0x0002 = Private
                | 0 /= w.&.0x0004 = Protected
                | otherwise       = Package

Add a code snippet to your website: www.paste.org