Stripe CTF Level 1 Solution

level1-miner.hs

module Main where

-- Solve Level 1 of Stripe
-- Create a sufficiently difficult git hash

-- ByteStrings are needed for hashing

import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as B8
import qualified Data.ByteString.Base16 as B16

-- This is stuff for manipulating Git
import qualified System.Process as P
import System.Exit (ExitCode(..))
import System.IO (withFile, IOMode(..), hGetContents)
import System.Environment (getArgs)
import System.FilePath ((</>))

-- Helps us modify LEDGER.txt
import Text.Regex (mkRegex, subRegex)

-- Parallel makes us faster
import Control.Parallel.Strategies (rseq,parMap)

-- so hash   such necessary   shamaze
import qualified Crypto.Hash.SHA1 as SHA

------------

-- Nothing complicated here; a simple type synonym
type CommitArgs = (String, String, String)

---- Making Git Hashes -----

-- Git object header
header :: B.ByteString -> B.ByteString
header body = B8.pack ("commit " ++ (show $ B.length body)) `B.append` (B.pack [0])

-- Add the git object header
addHeader :: B.ByteString -> B.ByteString
addHeader body = (header body) `B.append` body

-- Make a sha of a commit body
shaify :: String -> B.ByteString
shaify = SHA.hash . addHeader . B8.pack

-- Make a commit body
commitBody :: CommitArgs -> String
commitBody (t,p,n) = unlines [
  "tree " ++ t,
  "parent " ++ p,
  "author Name <email@domain.tld> " ++ n ++ " -0500",
  "committer Name <email@domain.tld> " ++ n ++ " -0500",
  "",
  "Give me a gitcoin"
  ]

-- Add an integer to to a commit body
commitCount :: String -> Int -> String
commitCount body cnt = body ++ (show cnt) ++ "\n"

-- Make an infinite list of commits that are different
commits :: CommitArgs -> [String]
commits args = let body = commitBody args
               in [body] ++ map (commitCount body) [1..]

-- Make an infinite list of commits that are different and their SHA1's
commitsShas :: CommitArgs -> [(String, B.ByteString)]
commitsShas = map (\b -> (b, shaify b)) . commits

-- Map an infinite list from commitsShas to one that attaches whether its difficult enough (in parallel)
mapDifficultyParallel :: B.ByteString -> [(String, B.ByteString)] -> [(Bool, (String, B.ByteString))]
mapDifficultyParallel diff arr = parMap rseq (\(s,bs) -> (bs < diff, (s, bs))) arr

-- Map an infinite list from commitsShas to one that attaches whether its difficult enough (serially)
mapDifficultySerial :: B.ByteString -> [(String, B.ByteString)] -> [(Bool, (String, B.ByteString))]
mapDifficultySerial diff arr = map (\(s,bs) -> (bs < diff, (s, bs))) arr

-- Filter out the first commit that passes the difficulty check
findDifficultCommit :: B.ByteString -> CommitArgs -> (String, B.ByteString)
findDifficultCommit diff = snd . head . filter fst . mapDifficultyParallel diff . commitsShas

----- Manipulating Git -----

readCommitArgs :: FilePath -> IO CommitArgs
readCommitArgs dir = do
  (_,tree,_) <- runGit dir ["write-tree"] ""
  (_,parent,_) <- runGit dir ["rev-parse", "HEAD"] ""
  (_,time,_) <- runCmd "date" ["+%s"] ""
  return (init tree, init parent, init time)

-- Helper to run a git command
runGit :: FilePath -> [String] -> String -> IO (ExitCode, String, String)
runGit dir args body = runCmd "git" (("--git-dir=" ++ dir </> ".git") : ("--exec-path=" ++ dir) : ("--work-tree=" ++ dir) : args) body

-- Helper to run a command
runCmd :: FilePath -> [String] -> String -> IO (ExitCode, String, String)
runCmd = P.readProcessWithExitCode

readDifficulty :: FilePath -> IO String
readDifficulty dir = readFile (dir </> "difficulty.txt")

-- Encodes a difficulty into a form so we can check against it
encodeDifficulty :: String -> B8.ByteString
encodeDifficulty = fst . B16.decode . B8.pack . init

-- Write a commit out using git hash-object
writeCommit :: FilePath -> String -> IO (ExitCode, String, String)
writeCommit dir commit = runGit dir ["hash-object", "-t", "commit", "--stdin", "-w"] commit

-- Reset to a commit sha and push to origin
resetAndPush :: FilePath -> String -> IO Bool
resetAndPush dir sha = do
  putStrLn $ "Resetting to " ++ sha
  runGit dir ["reset", "--hard", sha] ""
  putStrLn "Pushing to origin"
  (code,out,err) <- runGit dir ["push", "origin", "master"] ""
  putStrLn out
  putStrLn err
  return (if ExitSuccess == code then True else False)

-- Update the ledger file appropriately
updateLedgerFile :: FilePath -> String -> IO (ExitCode, String, String)
updateLedgerFile dir user = do
  let fn = dir </> "LEDGER.txt"
  nl <- withFile fn ReadMode (\h -> hGetContents h >>= \l -> putStrLn l >>= \_ -> return $ updateLedger user l)
  writeFile fn nl
  putStrLn "Updated ledger!"
  runGit dir ["add", "LEDGER.txt"] ""

-- Update the ledger contents
updateLedger :: String -> String -> String
updateLedger user ledger =
    case getCount copy of
      [] -> copy ++ user ++ ": 1\n"
      [i] -> subRegex (mkRegex $ user ++ ": " ++ (show i)) copy (user ++ ": " ++ (show $ i + 1))
      _ -> error "am i in there more than once?"
  where
    copy = map id ledger
    getCount = map snd . filter fst . map (\l -> (head (words l) == (user++":"), read (last (words l)) :: Int)) . lines

-- Reset the entire thing to start over (after a failed push)
reset :: FilePath -> IO ()
reset dir = do
  runGit dir ["fetch", "origin"] ""
  runGit dir ["reset", "--hard", "origin/master"] ""
  return ()

-- Main function to mine gitcoins
mineGitcoins :: FilePath -> String -> IO ()
mineGitcoins dir user = do
  reset dir
  updateLedgerFile dir user
  comargs <- readCommitArgs dir
  diff <- readDifficulty dir
  let bdiff = encodeDifficulty diff
  putStrLn $ show comargs
  let (commit, sha) = findDifficultCommit (bdiff) comargs
  let ssha = (B8.unpack . B16.encode) sha
  putStr commit
  putStrLn ssha
  writeCommit dir commit
  success <- resetAndPush dir ssha
  if success then putStrLn "SUCCESS!" else mineGitcoins dir user

-- Main function to parse args and run the miner
main :: IO ()
main = do
  args <- getArgs
  let dir = case args of [] -> "../one-up"
                         [d] -> d
  let user = case args of [_,u] -> u
                          _ -> "user-id"
  mineGitcoins dir user