module B9.DiskImages where
import B9.QCUtil
import Control.Parallel.Strategies
import Data.Binary
import Data.Data
import Data.Hashable
import Data.Map (Map)
import Data.Maybe
import Data.Semigroup as Sem
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Time.Calendar (Day (ModifiedJulianDay))
import Data.Time.Clock (UTCTime (..), diffTimeToPicoseconds, secondsToDiffTime)
import Data.Time.Format (defaultTimeLocale, formatTime, parseTimeOrError)
import GHC.Generics (Generic)
import System.FilePath
import Test.Hspec (Spec, describe, it)
import Test.QuickCheck
import qualified Text.PrettyPrint.Boxes as Boxes
import Text.Printf
data ImageTarget
= ImageTarget
ImageDestination
ImageSource
MountPoint
deriving (Read, Show, Typeable, Data, Eq, Generic)
instance Hashable ImageTarget
instance Binary ImageTarget
instance NFData ImageTarget
data MountPoint = MountPoint FilePath | NotMounted
deriving (Show, Read, Typeable, Data, Eq, Generic)
instance Hashable MountPoint
instance Binary MountPoint
instance NFData MountPoint
data ImageDestination
=
Share String ImageType ImageResize
|
LiveInstallerImage String FilePath ImageResize
|
LocalFile Image ImageResize
|
Transient
deriving (Read, Show, Typeable, Data, Eq, Generic)
instance Hashable ImageDestination
instance Binary ImageDestination
instance NFData ImageDestination
data ImageSource
=
EmptyImage String FileSystem ImageType ImageSize
|
CopyOnWrite Image
|
SourceImage Image Partition ImageResize
|
From String ImageResize
deriving (Show, Read, Typeable, Data, Eq, Generic)
instance Hashable ImageSource
instance Binary ImageSource
instance NFData ImageSource
data Partition
=
NoPT
|
Partition Int
deriving (Eq, Show, Read, Typeable, Data, Generic)
instance Hashable Partition
instance Binary Partition
instance NFData Partition
data Image = Image FilePath ImageType FileSystem
deriving (Eq, Show, Read, Typeable, Data, Generic)
instance Hashable Image
instance Binary Image
instance NFData Image
data ImageType = Raw | QCow2 | Vmdk
deriving (Eq, Read, Typeable, Data, Show, Generic)
instance CoArbitrary ImageType
instance Function ImageType
instance Hashable ImageType
instance Binary ImageType
instance NFData ImageType
data FileSystem = NoFileSystem | Ext4 | Ext4_64 | ISO9660 | VFAT
deriving (Eq, Show, Read, Typeable, Data, Generic)
instance Function FileSystem
instance CoArbitrary FileSystem
instance Hashable FileSystem
instance Binary FileSystem
instance NFData FileSystem
data ImageSize = ImageSize Int SizeUnit
deriving (Eq, Show, Read, Typeable, Data, Generic)
instance Hashable ImageSize
instance Binary ImageSize
instance NFData ImageSize
bytesToKiloBytes :: Int -> ImageSize
bytesToKiloBytes x =
let kbRoundedDown = x `div` 1024
rest = x `mod` 1024
kbRoundedUp = if rest > 0 then kbRoundedDown + 1 else kbRoundedDown
in ImageSize kbRoundedUp KB
imageSizeToKiB :: ImageSize -> Int
imageSizeToKiB (ImageSize size unit) =
size * sizeUnitKiB unit
sizeUnitKiB :: SizeUnit -> Int
sizeUnitKiB GB = 1024 * sizeUnitKiB MB
sizeUnitKiB MB = 1024 * sizeUnitKiB KB
sizeUnitKiB KB = 1
normalizeSize :: ImageSize -> ImageSize
normalizeSize i@(ImageSize _ GB) = i
normalizeSize i@(ImageSize size unit)
| size `mod` 1024 == 0 =
normalizeSize (ImageSize (size `div` 1024) (succ unit))
| otherwise = i
addImageSize :: ImageSize -> ImageSize -> ImageSize
addImageSize (ImageSize value unit) (ImageSize value' unit') =
normalizeSize
(ImageSize (value * sizeUnitKiB unit + value' * sizeUnitKiB unit') KB)
data SizeUnit = KB | MB | GB
deriving (Eq, Show, Read, Ord, Enum, Bounded, Typeable, Data, Generic)
instance Hashable SizeUnit
instance Binary SizeUnit
instance NFData SizeUnit
data ImageResize
=
ResizeImage ImageSize
|
Resize ImageSize
|
ShrinkToMinimumAndIncrease ImageSize
|
ShrinkToMinimum
|
KeepSize
deriving (Eq, Show, Read, Typeable, Data, Generic)
instance Hashable ImageResize
instance Binary ImageResize
instance NFData ImageResize
type Mounted a = (a, MountPoint)
data SharedImage
= SharedImage
SharedImageName
SharedImageDate
SharedImageBuildId
ImageType
FileSystem
deriving (Eq, Read, Show, Typeable, Data, Generic)
instance Arbitrary SharedImage where
arbitrary =
SharedImage
<$> smaller arbitrary
<*> smaller arbitrary
<*> smaller arbitrary
<*> smaller arbitrary
<*> smaller arbitrary
instance CoArbitrary SharedImage
instance Function SharedImage
instance Hashable SharedImage
instance Binary SharedImage
instance NFData SharedImage
newtype SharedImageName
= SharedImageName String
deriving (Eq, Ord, Read, Show, Typeable, Data, Hashable, Binary, NFData, CoArbitrary)
fromSharedImageName :: SharedImageName -> String
fromSharedImageName (SharedImageName b) = b
newtype SharedImageDate
= SharedImageDate String
deriving (Eq, Ord, Read, Show, Typeable, Data, Hashable, Binary, NFData, CoArbitrary)
newtype SharedImageBuildId
= SharedImageBuildId String
deriving
(Eq, Ord, Read, Show, Typeable, Data, Hashable, Binary, NFData, CoArbitrary)
fromSharedImageBuildId :: SharedImageBuildId -> String
fromSharedImageBuildId (SharedImageBuildId b) = b
instance Ord SharedImage where
compare (SharedImage n d b _ _) (SharedImage n' d' b' _ _) =
compare n n' Sem.<> compare d d' Sem.<> compare b b'
sharedImagesToMap :: [SharedImage] -> Map SharedImageName (Set SharedImage)
sharedImagesToMap _ = error "TODO IMPLEMENT ME"
takeLatestSharedImage :: [SharedImage] -> Maybe SharedImage
takeLatestSharedImage _ss = do
error "TODO IMPLEMENT ME"
imageFileName :: Image -> FilePath
imageFileName (Image f _ _) = f
imageImageType :: Image -> ImageType
imageImageType (Image _ t _) = t
getImageDestinationOutputFiles :: ImageTarget -> [FilePath]
getImageDestinationOutputFiles (ImageTarget d _ _) = case d of
LiveInstallerImage liName liPath _ ->
let path = liPath </> "machines" </> liName </> "disks" </> "raw"
in [path </> "0.raw", path </> "0.size", path </> "VERSION"]
LocalFile (Image lfPath _ _) _ -> [lfPath]
_ -> []
imageDestinationSharedImageName :: ImageDestination -> Maybe SharedImageName
imageDestinationSharedImageName (Share n _ _) = Just (SharedImageName n)
imageDestinationSharedImageName _ = Nothing
imageSourceSharedImageName :: ImageSource -> Maybe SharedImageName
imageSourceSharedImageName (From n _) = Just (SharedImageName n)
imageSourceSharedImageName _ = Nothing
itImageDestination :: ImageTarget -> ImageDestination
itImageDestination (ImageTarget d _ _) = d
itImageSource :: ImageTarget -> ImageSource
itImageSource (ImageTarget _ s _) = s
itImageMountPoint :: ImageTarget -> MountPoint
itImageMountPoint (ImageTarget _ _ m) = m
isPartitioned :: Partition -> Bool
isPartitioned p
| p == NoPT = False
| otherwise = True
getPartition :: Partition -> Int
getPartition (Partition p) = p
getPartition NoPT = error "No partitions!"
imageFileExtension :: ImageType -> String
imageFileExtension Raw = "raw"
imageFileExtension QCow2 = "qcow2"
imageFileExtension Vmdk = "vmdk"
changeImageFormat :: ImageType -> Image -> Image
changeImageFormat fmt' (Image img _ fs) = Image img' fmt' fs
where
img' = replaceExtension img (imageFileExtension fmt')
changeImageDirectory :: FilePath -> Image -> Image
changeImageDirectory dir (Image img fmt fs) = Image img' fmt fs
where
img' = dir </> takeFileName img
getImageSourceImageType :: ImageSource -> Maybe ImageType
getImageSourceImageType (EmptyImage _ _ t _) = Just t
getImageSourceImageType (CopyOnWrite i) = Just $ imageImageType i
getImageSourceImageType (SourceImage i _ _) = Just $ imageImageType i
getImageSourceImageType (From _ _) = Nothing
sharedImageName :: SharedImage -> SharedImageName
sharedImageName (SharedImage n _ _ _ _) = n
sharedImageDate :: SharedImage -> SharedImageDate
sharedImageDate (SharedImage _ n _ _ _) = n
sharedImageBuildId :: SharedImage -> SharedImageBuildId
sharedImageBuildId (SharedImage _ _ n _ _) = n
prettyPrintSharedImages :: Set SharedImage -> String
prettyPrintSharedImages imgs = Boxes.render table
where
table = Boxes.hsep 1 Boxes.left cols
where
cols = [nameC, dateC, idC]
where
nameC = col "Name" ((\(SharedImageName n) -> n) . sharedImageName)
dateC = col "Date" ((\(SharedImageDate n) -> n) . sharedImageDate)
idC =
col
"ID"
((\(SharedImageBuildId n) -> n) . sharedImageBuildId)
col title accessor =
Boxes.text title Boxes.// Boxes.vcat Boxes.left cells
where
cells = Boxes.text . accessor <$> Set.toList imgs
sharedImageImage :: SharedImage -> Image
sharedImageImage (SharedImage (SharedImageName n) _ (SharedImageBuildId bid) sharedImageType sharedImageFileSystem) =
Image
(n ++ "_" ++ bid <.> imageFileExtension sharedImageType)
sharedImageType
sharedImageFileSystem
sharedImageFileName :: SharedImage -> FilePath
sharedImageFileName (SharedImage (SharedImageName n) _ (SharedImageBuildId bid) _ _) =
n ++ "_" ++ bid <.> sharedImageFileExtension
sharedImagesRootDirectory :: FilePath
sharedImagesRootDirectory = "b9_shared_images"
sharedImageFileExtension :: String
sharedImageFileExtension = "b9si"
sharedImageDefaultImageType :: ImageType
sharedImageDefaultImageType = QCow2
transientCOWImage :: FilePath -> FilePath -> ImageTarget
transientCOWImage fileName mountPoint =
ImageTarget
Transient
(CopyOnWrite (Image fileName QCow2 Ext4))
(MountPoint mountPoint)
transientSharedImage :: SharedImageName -> FilePath -> ImageTarget
transientSharedImage (SharedImageName name) mountPoint =
ImageTarget Transient (From name KeepSize) (MountPoint mountPoint)
transientLocalImage :: FilePath -> FilePath -> ImageTarget
transientLocalImage name mountPoint =
ImageTarget Transient (From name KeepSize) (MountPoint mountPoint)
shareCOWImage :: FilePath -> SharedImageName -> FilePath -> ImageTarget
shareCOWImage srcFilename (SharedImageName destName) mountPoint =
ImageTarget
(Share destName QCow2 KeepSize)
(CopyOnWrite (Image srcFilename QCow2 Ext4))
(MountPoint mountPoint)
shareSharedImage ::
SharedImageName -> SharedImageName -> FilePath -> ImageTarget
shareSharedImage (SharedImageName srcName) (SharedImageName destName) mountPoint =
ImageTarget
(Share destName QCow2 KeepSize)
(From srcName KeepSize)
(MountPoint mountPoint)
shareLocalImage :: FilePath -> SharedImageName -> FilePath -> ImageTarget
shareLocalImage srcName (SharedImageName destName) mountPoint =
ImageTarget
(Share destName QCow2 KeepSize)
(SourceImage (Image srcName QCow2 Ext4) NoPT KeepSize)
(MountPoint mountPoint)
cowToliveInstallerImage ::
String -> FilePath -> FilePath -> FilePath -> ImageTarget
cowToliveInstallerImage srcName destName outDir mountPoint =
ImageTarget
(LiveInstallerImage destName outDir KeepSize)
(CopyOnWrite (Image srcName QCow2 Ext4))
(MountPoint mountPoint)
cowToLocalImage :: FilePath -> FilePath -> FilePath -> ImageTarget
cowToLocalImage srcName destName mountPoint =
ImageTarget
(LocalFile (Image destName QCow2 Ext4) KeepSize)
(CopyOnWrite (Image srcName QCow2 Ext4))
(MountPoint mountPoint)
localToLocalImage :: FilePath -> FilePath -> FilePath -> ImageTarget
localToLocalImage srcName destName mountPoint =
ImageTarget
(LocalFile (Image destName QCow2 Ext4) KeepSize)
(SourceImage (Image srcName QCow2 Ext4) NoPT KeepSize)
(MountPoint mountPoint)
partition1ToLocalImage :: FilePath -> FilePath -> FilePath -> ImageTarget
partition1ToLocalImage srcName destName mountPoint =
ImageTarget
(LocalFile (Image destName QCow2 Ext4) KeepSize)
(SourceImage (Image srcName QCow2 Ext4) NoPT KeepSize)
(MountPoint mountPoint)
splitToIntermediateSharedImage ::
ImageTarget -> SharedImageName -> (ImageTarget, ImageTarget)
splitToIntermediateSharedImage (ImageTarget dst src mnt) (SharedImageName intermediateName) =
(imgTargetShared, imgTargetExport)
where
imgTargetShared = ImageTarget intermediateTo src mnt
imgTargetExport = ImageTarget dst intermediateFrom mnt
intermediateTo =
Share
intermediateName
(fromMaybe sharedImageDefaultImageType (getImageSourceImageType src))
KeepSize
intermediateFrom = From intermediateName KeepSize
instance Arbitrary ImageTarget where
arbitrary =
ImageTarget
<$> smaller arbitrary
<*> smaller arbitrary
<*> smaller arbitrary
instance Arbitrary ImageSource where
arbitrary =
oneof
[ EmptyImage . printf "img-label-%0X"
<$> choose (0, 63 :: Int)
<*> smaller arbitrary
<*> smaller arbitrary
<*> smaller arbitrary,
CopyOnWrite <$> smaller arbitrary,
SourceImage
<$> smaller arbitrary
<*> smaller arbitrary
<*> smaller arbitrary,
From <$> arbitrarySharedImageName <*> smaller arbitrary
]
instance Arbitrary ImageDestination where
arbitrary =
oneof
[ Share
<$> arbitrarySharedImageName
<*> smaller arbitrary
<*> smaller arbitrary,
LiveInstallerImage "live-installer" "output-path"
<$> smaller arbitrary,
pure Transient
]
instance Arbitrary MountPoint where
arbitrary = elements [MountPoint "/mnt", NotMounted]
instance Arbitrary ImageResize where
arbitrary =
oneof
[ ResizeImage <$> smaller arbitrary,
Resize <$> smaller arbitrary,
ShrinkToMinimumAndIncrease <$> smaller arbitrary,
pure ShrinkToMinimum,
pure KeepSize
]
instance Arbitrary Partition where
arbitrary = oneof [Partition <$> choose (0, 2), pure NoPT]
instance Arbitrary Image where
arbitrary =
Image . printf "img-file-name-%0X"
<$> choose (0, 63 :: Int)
<*> smaller arbitrary
<*> smaller arbitrary
instance Arbitrary FileSystem where
arbitrary = elements [Ext4]
instance Arbitrary ImageType where
arbitrary = elements [Raw, QCow2, Vmdk]
instance Arbitrary ImageSize where
arbitrary = ImageSize <$> smaller arbitrary <*> smaller arbitrary
instance Arbitrary SizeUnit where
arbitrary = elements [KB, MB, GB]
instance Arbitrary SharedImageName where
arbitrary = SharedImageName <$> arbitrarySharedImageName
instance Function SharedImageName where
function = functionShow
arbitrarySharedImageName :: Gen String
arbitrarySharedImageName =
printf "shared-img-%0X" <$> choose (0, 63 :: Int)
instance Arbitrary SharedImageBuildId where
arbitrary = do
SharedImageBuildId . printf "shared-img-build-id-%0X" <$> choose (0, maxBound `div` 1024 :: Int)
instance Function SharedImageBuildId where
function = functionMap fromSharedImageBuildId SharedImageBuildId
instance Function SharedImageDate where
function =
functionMap
( ( \(UTCTime (ModifiedJulianDay d) dt) ->
(d, diffTimeToPicoseconds dt `div` 1000000000000)
)
. parseTimeOrError False defaultTimeLocale "%F-%T"
. (\(SharedImageDate d) -> d)
)
( SharedImageDate
. (formatTime defaultTimeLocale "%F-%T")
. (\(d, dt) -> UTCTime (ModifiedJulianDay d) (secondsToDiffTime dt))
)
instance Arbitrary SharedImageDate where
arbitrary =
SharedImageDate
. (formatTime defaultTimeLocale "%F-%T")
. (\(d, dt) -> UTCTime (ModifiedJulianDay d) (secondsToDiffTime dt))
<$> arbitrary
unitTests :: Spec
unitTests =
describe "ImageSize"
$ describe "bytesToKiloBytes"
$ do
it "accepts maxBound" $
toInteger (imageSizeToKiB (bytesToKiloBytes maxBound)) * 1024 === toInteger (maxBound :: Int) + 1
it "doesn't decrease in size" $
property
( \(x :: Int) ->
x <= maxBound - 1024 ==> label "bytesToKiloBytes x >= x" (imageSizeToKiB (bytesToKiloBytes x) >= (x `div` 1024))
)