-- | Utility module to extract a primary partition from an MBR partition on a
--    raw image file.
module B9.MBR
  ( getPartition,
    PrimaryPartition (..),
    MBR (..),
    CHS (..),
  )
where

import Data.Binary.Get
import qualified Data.ByteString.Lazy as BL
import Data.Word
import Text.Printf

getPartition :: Int -> FilePath -> IO (Word64, Word64)
getPartition :: Int -> FilePath -> IO (Word64, Word64)
getPartition Int
n FilePath
f = ByteString -> (Word64, Word64)
decodeMBR (ByteString -> (Word64, Word64))
-> IO ByteString -> IO (Word64, Word64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> IO ByteString
BL.readFile FilePath
f
  where
    decodeMBR :: ByteString -> (Word64, Word64)
decodeMBR ByteString
input =
      let mbr :: MBR
mbr = Get MBR -> ByteString -> MBR
forall a. Get a -> ByteString -> a
runGet Get MBR
getMBR ByteString
input
          part :: PrimaryPartition
part =
            ( case Int
n of
                Int
1 -> MBR -> PrimaryPartition
mbrPart1
                Int
2 -> MBR -> PrimaryPartition
mbrPart2
                Int
3 -> MBR -> PrimaryPartition
mbrPart3
                Int
4 -> MBR -> PrimaryPartition
mbrPart4
                Int
b ->
                  FilePath -> MBR -> PrimaryPartition
forall a. HasCallStack => FilePath -> a
error
                    ( FilePath -> Int -> FilePath -> FilePath
forall r. PrintfType r => FilePath -> r
printf
                        FilePath
"Error: Invalid partition index %i only partitions 1-4 are allowed. Image file: '%s'"
                        Int
b
                        FilePath
f
                    )
            )
              MBR
mbr
          start :: Word64
start = Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PrimaryPartition -> Word32
primPartLbaStart PrimaryPartition
part)
          len :: Word64
len = Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PrimaryPartition -> Word32
primPartSectors PrimaryPartition
part)
       in (Word64
start Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
sectorSize, Word64
len Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
* Word64
sectorSize)

sectorSize :: Word64
sectorSize :: Word64
sectorSize = Word64
512

bootCodeSize :: Int
bootCodeSize :: Int
bootCodeSize = Int
446

data MBR
  = MBR
      { MBR -> PrimaryPartition
mbrPart1 :: !PrimaryPartition,
        MBR -> PrimaryPartition
mbrPart2 :: !PrimaryPartition,
        MBR -> PrimaryPartition
mbrPart3 :: !PrimaryPartition,
        MBR -> PrimaryPartition
mbrPart4 :: !PrimaryPartition
      }
  deriving (Int -> MBR -> FilePath -> FilePath
[MBR] -> FilePath -> FilePath
MBR -> FilePath
(Int -> MBR -> FilePath -> FilePath)
-> (MBR -> FilePath) -> ([MBR] -> FilePath -> FilePath) -> Show MBR
forall a.
(Int -> a -> FilePath -> FilePath)
-> (a -> FilePath) -> ([a] -> FilePath -> FilePath) -> Show a
showList :: [MBR] -> FilePath -> FilePath
$cshowList :: [MBR] -> FilePath -> FilePath
show :: MBR -> FilePath
$cshow :: MBR -> FilePath
showsPrec :: Int -> MBR -> FilePath -> FilePath
$cshowsPrec :: Int -> MBR -> FilePath -> FilePath
Show)

data PrimaryPartition
  = PrimaryPartition
      { PrimaryPartition -> Word8
primPartStatus :: !Word8,
        PrimaryPartition -> CHS
primPartChsStart :: !CHS,
        PrimaryPartition -> Word8
primPartPartType :: !Word8,
        PrimaryPartition -> CHS
primPartChsEnd :: !CHS,
        PrimaryPartition -> Word32
primPartLbaStart :: !Word32,
        PrimaryPartition -> Word32
primPartSectors :: !Word32
      }
  deriving (Int -> PrimaryPartition -> FilePath -> FilePath
[PrimaryPartition] -> FilePath -> FilePath
PrimaryPartition -> FilePath
(Int -> PrimaryPartition -> FilePath -> FilePath)
-> (PrimaryPartition -> FilePath)
-> ([PrimaryPartition] -> FilePath -> FilePath)
-> Show PrimaryPartition
forall a.
(Int -> a -> FilePath -> FilePath)
-> (a -> FilePath) -> ([a] -> FilePath -> FilePath) -> Show a
showList :: [PrimaryPartition] -> FilePath -> FilePath
$cshowList :: [PrimaryPartition] -> FilePath -> FilePath
show :: PrimaryPartition -> FilePath
$cshow :: PrimaryPartition -> FilePath
showsPrec :: Int -> PrimaryPartition -> FilePath -> FilePath
$cshowsPrec :: Int -> PrimaryPartition -> FilePath -> FilePath
Show)

data CHS
  = CHS
      { CHS -> Word8
chsH :: !Word8,
        CHS -> Word8
chs_CUpper2_S :: !Word8,
        CHS -> Word8
chs_CLower8 :: !Word8
      }
  deriving (Int -> CHS -> FilePath -> FilePath
[CHS] -> FilePath -> FilePath
CHS -> FilePath
(Int -> CHS -> FilePath -> FilePath)
-> (CHS -> FilePath) -> ([CHS] -> FilePath -> FilePath) -> Show CHS
forall a.
(Int -> a -> FilePath -> FilePath)
-> (a -> FilePath) -> ([a] -> FilePath -> FilePath) -> Show a
showList :: [CHS] -> FilePath -> FilePath
$cshowList :: [CHS] -> FilePath -> FilePath
show :: CHS -> FilePath
$cshow :: CHS -> FilePath
showsPrec :: Int -> CHS -> FilePath -> FilePath
$cshowsPrec :: Int -> CHS -> FilePath -> FilePath
Show)

getMBR :: Get MBR
getMBR :: Get MBR
getMBR =
  Int -> Get ()
skip Int
bootCodeSize Get () -> Get MBR -> Get MBR
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> PrimaryPartition
-> PrimaryPartition -> PrimaryPartition -> PrimaryPartition -> MBR
MBR (PrimaryPartition
 -> PrimaryPartition -> PrimaryPartition -> PrimaryPartition -> MBR)
-> Get PrimaryPartition
-> Get
     (PrimaryPartition -> PrimaryPartition -> PrimaryPartition -> MBR)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get PrimaryPartition
getPart Get
  (PrimaryPartition -> PrimaryPartition -> PrimaryPartition -> MBR)
-> Get PrimaryPartition
-> Get (PrimaryPartition -> PrimaryPartition -> MBR)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get PrimaryPartition
getPart Get (PrimaryPartition -> PrimaryPartition -> MBR)
-> Get PrimaryPartition -> Get (PrimaryPartition -> MBR)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get PrimaryPartition
getPart Get (PrimaryPartition -> MBR) -> Get PrimaryPartition -> Get MBR
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get PrimaryPartition
getPart

getPart :: Get PrimaryPartition
getPart :: Get PrimaryPartition
getPart =
  Word8
-> CHS -> Word8 -> CHS -> Word32 -> Word32 -> PrimaryPartition
PrimaryPartition
    (Word8
 -> CHS -> Word8 -> CHS -> Word32 -> Word32 -> PrimaryPartition)
-> Get Word8
-> Get
     (CHS -> Word8 -> CHS -> Word32 -> Word32 -> PrimaryPartition)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
getWord8
    Get (CHS -> Word8 -> CHS -> Word32 -> Word32 -> PrimaryPartition)
-> Get CHS
-> Get (Word8 -> CHS -> Word32 -> Word32 -> PrimaryPartition)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get CHS
getCHS
    Get (Word8 -> CHS -> Word32 -> Word32 -> PrimaryPartition)
-> Get Word8 -> Get (CHS -> Word32 -> Word32 -> PrimaryPartition)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word8
getWord8
    Get (CHS -> Word32 -> Word32 -> PrimaryPartition)
-> Get CHS -> Get (Word32 -> Word32 -> PrimaryPartition)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get CHS
getCHS
    Get (Word32 -> Word32 -> PrimaryPartition)
-> Get Word32 -> Get (Word32 -> PrimaryPartition)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word32
getWord32le
    Get (Word32 -> PrimaryPartition)
-> Get Word32 -> Get PrimaryPartition
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word32
getWord32le

getCHS :: Get CHS
getCHS :: Get CHS
getCHS = Word8 -> Word8 -> Word8 -> CHS
CHS (Word8 -> Word8 -> Word8 -> CHS)
-> Get Word8 -> Get (Word8 -> Word8 -> CHS)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
getWord8 Get (Word8 -> Word8 -> CHS) -> Get Word8 -> Get (Word8 -> CHS)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word8
getWord8 Get (Word8 -> CHS) -> Get Word8 -> Get CHS
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Word8
getWord8