{-# LANGUAGE FlexibleContexts #-}

module Propellor.Property.Parted (
	-- * Types
	TableType(..),
	PartTable(..),
	partTableSize,
	Partition(..),
	mkPartition,
	Partition.Fs(..),
	PartSize(..),
	ByteSize,
	toPartSize,
	fromPartSize,
	reducePartSize,
	Alignment(..),
	safeAlignment,
	Partition.MkfsOpts,
	PartType(..),
	PartFlag(..),
	-- * Properties
	partitioned,
	parted,
	Eep(..),
	installed,
	-- * Partition table sizing
	calcPartTable,
	DiskSize(..),
	DiskPart,
	DiskSpaceUse(..),
	useDiskSpace,
	defSz,
	fudgeSz,
) where

import Propellor.Base
import Propellor.Property.Parted.Types
import qualified Propellor.Property.Apt as Apt
import qualified Propellor.Property.Pacman as Pacman
import qualified Propellor.Property.Partition as Partition
import Propellor.Types.PartSpec (PartSpec)
import Utility.DataUnits

import System.Posix.Files
import qualified Data.Semigroup as Sem
import Data.List (genericLength)

data Eep = YesReallyDeleteDiskContents

-- | Partitions a disk using parted, and formats the partitions.
--
-- The FilePath can be a block device (eg, \/dev\/sda), or a disk image file.
--
-- This deletes any existing partitions in the disk! Use with EXTREME caution!
partitioned :: Eep -> FilePath -> PartTable -> Property DebianLike
partitioned :: Eep
-> FilePath
-> PartTable
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish])
partitioned Eep
eep FilePath
disk parttable :: PartTable
parttable@(PartTable TableType
_ Alignment
_ [Partition]
parts) = forall {k} (metatypes :: k).
SingI metatypes =>
FilePath
-> (OuterMetaTypesWitness metatypes -> Propellor Result)
-> Property (MetaTypes metatypes)
property' FilePath
desc forall a b. (a -> b) -> a -> b
$ \OuterMetaTypesWitness
  '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish]
w -> do
	Bool
isdev <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ FileStatus -> Bool
isBlockDevice forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> IO FileStatus
getFileStatus FilePath
disk
	forall (inner :: [MetaType]) (outer :: [MetaType]).
EnsurePropertyAllowed inner outer =>
OuterMetaTypesWitness outer
-> Property (MetaTypes inner) -> Propellor Result
ensureProperty OuterMetaTypesWitness
  '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish]
w forall a b. (a -> b) -> a -> b
$ forall {k} (metatypes :: k).
SingI metatypes =>
FilePath
-> Props (MetaTypes metatypes) -> Property (MetaTypes metatypes)
combineProperties FilePath
desc forall a b. (a -> b) -> a -> b
$ Props UnixLike
props
		forall {a} p (y :: [a]) (x :: [a]).
(IsProp p, MetaTypes y ~ GetMetaTypes p,
 CheckCombinableNote x y (NoteFor ('Text "&"))) =>
Props (MetaTypes x) -> p -> Props (MetaTypes (Combine x y))
& Eep
-> FilePath
-> [FilePath]
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish]
      + ArchLinux)
parted Eep
eep FilePath
disk (forall a b. (a, b) -> a
fst (PartTable -> ([FilePath], ByteSize)
calcPartedParamsSize PartTable
parttable))
		forall {a} p (y :: [a]) (x :: [a]).
(IsProp p, MetaTypes y ~ GetMetaTypes p,
 CheckCombinableNote x y (NoteFor ('Text "&"))) =>
Props (MetaTypes x) -> p -> Props (MetaTypes (Combine x y))
& if Bool
isdev
			then [FilePath]
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish])
formatl (forall a b. (a -> b) -> [a] -> [b]
map (\Int
n -> FilePath
disk forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> FilePath
show Int
n) [Int
1 :: Int ..])
			else FilePath
-> ([LoopDev]
    -> Property
         (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish]))
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish])
Partition.kpartx FilePath
disk ([FilePath]
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish])
formatl forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map LoopDev -> FilePath
Partition.partitionLoopDev)
  where
	desc :: FilePath
desc = FilePath
disk forall a. [a] -> [a] -> [a]
++ FilePath
" partitioned"
	formatl :: [FilePath]
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish])
formatl [FilePath]
devs = forall {k} (metatypes :: k).
SingI metatypes =>
FilePath
-> Props (MetaTypes metatypes) -> Property (MetaTypes metatypes)
combineProperties FilePath
desc (forall {k} (metatypes :: k).
[Property (MetaTypes metatypes)] -> Props (MetaTypes metatypes)
toProps forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Partition, FilePath)
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish])
format (forall a b. [a] -> [b] -> [(a, b)]
zip [Partition]
parts [FilePath]
devs))
	format :: (Partition, FilePath)
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish])
format (Partition
p, FilePath
dev) = case Partition -> Maybe Fs
partFs Partition
p of
		Just Fs
fs -> [FilePath]
-> Eep
-> Fs
-> FilePath
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish])
Partition.formatted' (Partition -> [FilePath]
partMkFsOpts Partition
p)
			Eep
Partition.YesReallyFormatPartition Fs
fs FilePath
dev
		Maybe Fs
Nothing -> forall {k} (t :: k). SingI t => Property (MetaTypes t)
doNothing

-- | Gets the total size of the disk specified by the partition table.
partTableSize :: PartTable -> ByteSize
partTableSize :: PartTable -> ByteSize
partTableSize = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. PartTable -> ([FilePath], ByteSize)
calcPartedParamsSize

calcPartedParamsSize :: PartTable -> ([String], ByteSize)
calcPartedParamsSize :: PartTable -> ([FilePath], ByteSize)
calcPartedParamsSize (PartTable TableType
tabletype Alignment
alignment [Partition]
parts) = 
	let ([[FilePath]]
ps, ByteSize
sz) = forall {a}.
(Num a, Show a) =>
a
-> ByteSize
-> [Partition]
-> [[FilePath]]
-> ([[FilePath]], ByteSize)
calcparts (ByteSize
1 :: Integer) ByteSize
firstpos [Partition]
parts []
	in (forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([FilePath]
mklabel forall a. a -> [a] -> [a]
: [[FilePath]]
ps), ByteSize
sz)
  where
	mklabel :: [FilePath]
mklabel = [FilePath
"mklabel", forall a. PartedVal a => a -> FilePath
pval TableType
tabletype]
	mkflag :: a -> (a, a) -> [FilePath]
mkflag a
partnum (a
f, a
b) =
		[ FilePath
"set"
		, forall a. Show a => a -> FilePath
show a
partnum
		, forall a. PartedVal a => a -> FilePath
pval a
f
		, forall a. PartedVal a => a -> FilePath
pval a
b
		]
	mkpart :: a -> a -> a -> Partition -> [FilePath]
mkpart a
partnum a
startpos a
endpos Partition
p = forall a. [Maybe a] -> [a]
catMaybes
		[ forall a. a -> Maybe a
Just FilePath
"mkpart"
		, forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. PartedVal a => a -> FilePath
pval (Partition -> PartType
partType Partition
p)
		, forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. PartedVal a => a -> FilePath
pval (Partition -> Maybe Fs
partFs Partition
p)
		, forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {a}. (Ord a, Num a, Show a) => a -> FilePath
partposexact a
startpos
		, forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {a}. Integral a => a -> FilePath
partposfuzzy a
endpos
		] forall a. [a] -> [a] -> [a]
++ case Partition -> Maybe FilePath
partName Partition
p of
			Just FilePath
n -> [FilePath
"name", forall a. Show a => a -> FilePath
show a
partnum, FilePath
n]
			Maybe FilePath
Nothing -> []
	calcparts :: a
-> ByteSize
-> [Partition]
-> [[FilePath]]
-> ([[FilePath]], ByteSize)
calcparts a
partnum ByteSize
startpos (Partition
p:[Partition]
ps) [[FilePath]]
c =
		let endpos :: ByteSize
endpos = ByteSize
startpos forall a. Num a => a -> a -> a
+ PartSize -> ByteSize
align (Partition -> PartSize
partSize Partition
p)
		in a
-> ByteSize
-> [Partition]
-> [[FilePath]]
-> ([[FilePath]], ByteSize)
calcparts (a
partnumforall a. Num a => a -> a -> a
+a
1) ByteSize
endpos [Partition]
ps
			([[FilePath]]
c forall a. [a] -> [a] -> [a]
++ forall {a} {a} {a}.
(Integral a, Show a, Show a, Num a, Ord a) =>
a -> a -> a -> Partition -> [FilePath]
mkpart a
partnum ByteSize
startpos (ByteSize
endposforall a. Num a => a -> a -> a
-ByteSize
1) Partition
p forall a. a -> [a] -> [a]
: forall a b. (a -> b) -> [a] -> [b]
map (forall {a} {a} {a}.
(Show a, PartedVal a, PartedVal a) =>
a -> (a, a) -> [FilePath]
mkflag a
partnum) (Partition -> [(PartFlag, Bool)]
partFlags Partition
p))
	calcparts a
_ ByteSize
endpos [] [[FilePath]]
c = ([[FilePath]]
c, ByteSize
endpos)

	-- Exact partition position value for parted.
	-- For alignment to work, the start of a partition must be
	-- specified exactly.
	partposexact :: a -> FilePath
partposexact a
n
		| a
n forall a. Ord a => a -> a -> Bool
> a
0 = forall a. Show a => a -> FilePath
show a
n forall a. [a] -> [a] -> [a]
++ FilePath
"B"
		-- parted can't make partitions smaller than 1MB;
		-- avoid failure in edge cases
		| Bool
otherwise = FilePath
"1MB"
	
	-- Fuzzy partition position valie for parted.
	-- This is used to specify the end of the partition,
	-- parted takes the "MB" as license to slightly reduce the
	-- partition size when something about the partition table
	-- does not allow the partition to end exactly at the position.
	partposfuzzy :: a -> FilePath
partposfuzzy a
n
		| a
n forall a. Ord a => a -> a -> Bool
> a
0 = forall a. Show a => a -> FilePath
show (forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n forall a. Fractional a => a -> a -> a
/ Double
1000000 :: Double) forall a. [a] -> [a] -> [a]
++ FilePath
"MB"
		| Bool
otherwise = FilePath
"1MB"

	-- Location of the start of the first partition,
	-- leaving space for the partition table, and aligning.
	firstpos :: ByteSize
firstpos = PartSize -> ByteSize
align PartSize
partitionTableOverhead
	
	align :: PartSize -> ByteSize
align = Alignment -> PartSize -> ByteSize
alignTo Alignment
alignment

-- | Runs parted on a disk with the specified parameters.
--
-- Parted is run in script mode, so it will never prompt for input.
parted :: Eep -> FilePath -> [String] -> Property (DebianLike + ArchLinux)
parted :: Eep
-> FilePath
-> [FilePath]
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish]
      + ArchLinux)
parted Eep
YesReallyDeleteDiskContents FilePath
disk [FilePath]
ps = Property UnixLike
p forall x y. Combines x y => x -> y -> CombinedType x y
`requires` Property
  (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish]
   + ArchLinux)
installed
  where
	p :: Property UnixLike
p = FilePath -> [FilePath] -> UncheckedProperty UnixLike
cmdProperty FilePath
"parted" (FilePath
"--script"forall a. a -> [a] -> [a]
:FilePath
"--align"forall a. a -> [a] -> [a]
:FilePath
"none"forall a. a -> [a] -> [a]
:FilePath
diskforall a. a -> [a] -> [a]
:[FilePath]
ps)
		forall (p :: * -> *) i.
Checkable p i =>
p i -> Result -> Property i
`assume` Result
MadeChange

-- | Gets parted installed.
installed :: Property (DebianLike + ArchLinux)
installed :: Property
  (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish]
   + ArchLinux)
installed = [FilePath]
-> Property
     (MetaTypes '[ 'Targeting 'OSDebian, 'Targeting 'OSBuntish])
Apt.installed [FilePath
"parted"] forall {k} ka kb (c :: k) (a :: ka) (b :: kb).
(HasCallStack, SingKind 'KProxy, SingKind 'KProxy,
 DemoteRep 'KProxy ~ [MetaType], DemoteRep 'KProxy ~ [MetaType],
 SingI c) =>
Property (MetaTypes a)
-> Property (MetaTypes b) -> Property (MetaTypes c)
`pickOS` [FilePath] -> Property ArchLinux
Pacman.installed [FilePath
"parted"]

-- | Some disk is used to store the partition table itself. Assume less
-- than 1 mb.
partitionTableOverhead :: PartSize
partitionTableOverhead :: PartSize
partitionTableOverhead = ByteSize -> PartSize
MegaBytes ByteSize
1

-- | Calculate a partition table, for a given size of disk.
--
-- For example:
--
-- >	calcPartTable (DiskSize (1024 * 1024 * 1024 * 100)) MSDOS safeAlignment
-- > 		[ partition EXT2 `mountedAt` "/boot"
-- > 			`setSize` MegaBytes 256
-- > 			`setFlag` BootFlag
-- >		, partition EXT4 `mountedAt` "/"
-- >			`useDiskSpace` RemainingSpace
-- >		]
calcPartTable :: DiskSize -> TableType -> Alignment -> [PartSpec DiskPart] -> PartTable
calcPartTable :: DiskSize
-> TableType -> Alignment -> [PartSpec DiskPart] -> PartTable
calcPartTable (DiskSize ByteSize
disksize) TableType
tt Alignment
alignment [PartSpec DiskPart]
l =
	TableType -> Alignment -> [Partition] -> PartTable
PartTable TableType
tt Alignment
alignment (forall a b. (a -> b) -> [a] -> [b]
map PartSpec DiskPart -> Partition
go [PartSpec DiskPart]
l)
  where
	go :: PartSpec DiskPart -> Partition
go (Maybe FilePath
_, MountOpts
_, PartSize -> Partition
mkpart, DiskPart
FixedDiskPart) = PartSize -> Partition
mkpart PartSize
defSz
	go (Maybe FilePath
_, MountOpts
_, PartSize -> Partition
mkpart, DynamicDiskPart (Percent Int
p)) = PartSize -> Partition
mkpart forall a b. (a -> b) -> a -> b
$ ByteSize -> PartSize
Bytes forall a b. (a -> b) -> a -> b
$
		ByteSize
diskremainingafterfixed forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
p forall a. Integral a => a -> a -> a
`div` ByteSize
100
	go (Maybe FilePath
_, MountOpts
_, PartSize -> Partition
mkpart, DynamicDiskPart DiskSpaceUse
RemainingSpace) = PartSize -> Partition
mkpart forall a b. (a -> b) -> a -> b
$ ByteSize -> PartSize
Bytes forall a b. (a -> b) -> a -> b
$
		ByteSize
diskremaining forall a. Integral a => a -> a -> a
`div` forall i a. Num i => [a] -> i
genericLength (forall a. (a -> Bool) -> [a] -> [a]
filter forall {a} {b} {c}. (a, b, c, DiskPart) -> Bool
isremainingspace [PartSpec DiskPart]
l)
	diskremainingafterfixed :: ByteSize
diskremainingafterfixed =
		ByteSize
disksize forall a. Num a => a -> a -> a
- [PartSpec DiskPart] -> ByteSize
sumsizes (forall a. (a -> Bool) -> [a] -> [a]
filter forall {a} {b} {c}. (a, b, c, DiskPart) -> Bool
isfixed [PartSpec DiskPart]
l)
	diskremaining :: ByteSize
diskremaining =
		ByteSize
disksize forall a. Num a => a -> a -> a
- [PartSpec DiskPart] -> ByteSize
sumsizes (forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {a} {b} {c}. (a, b, c, DiskPart) -> Bool
isremainingspace) [PartSpec DiskPart]
l)
	sumsizes :: [PartSpec DiskPart] -> ByteSize
sumsizes = PartTable -> ByteSize
partTableSize forall b c a. (b -> c) -> (a -> b) -> a -> c
. TableType -> Alignment -> [Partition] -> PartTable
PartTable TableType
tt Alignment
alignment forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map PartSpec DiskPart -> Partition
go
	isfixed :: (a, b, c, DiskPart) -> Bool
isfixed (a
_, b
_, c
_, DiskPart
FixedDiskPart) = Bool
True
	isfixed (a, b, c, DiskPart)
_ = Bool
False
	isremainingspace :: (a, b, c, DiskPart) -> Bool
isremainingspace (a
_, b
_, c
_, DynamicDiskPart DiskSpaceUse
RemainingSpace) = Bool
True
	isremainingspace (a, b, c, DiskPart)
_ = Bool
False

-- | Size of a disk, in bytes.
newtype DiskSize = DiskSize ByteSize
	deriving (Int -> DiskSize -> ShowS
[DiskSize] -> ShowS
DiskSize -> FilePath
forall a.
(Int -> a -> ShowS) -> (a -> FilePath) -> ([a] -> ShowS) -> Show a
showList :: [DiskSize] -> ShowS
$cshowList :: [DiskSize] -> ShowS
show :: DiskSize -> FilePath
$cshow :: DiskSize -> FilePath
showsPrec :: Int -> DiskSize -> ShowS
$cshowsPrec :: Int -> DiskSize -> ShowS
Show)

data DiskPart = FixedDiskPart | DynamicDiskPart DiskSpaceUse

data DiskSpaceUse = Percent Int | RemainingSpace

instance Sem.Semigroup DiskPart where
	DiskPart
FixedDiskPart <> :: DiskPart -> DiskPart -> DiskPart
<> DiskPart
FixedDiskPart = DiskPart
FixedDiskPart
	DynamicDiskPart (Percent Int
a) <> DynamicDiskPart (Percent Int
b) =
		DiskSpaceUse -> DiskPart
DynamicDiskPart (Int -> DiskSpaceUse
Percent (Int
a forall a. Num a => a -> a -> a
+ Int
b))
	DynamicDiskPart DiskSpaceUse
RemainingSpace <> DynamicDiskPart DiskSpaceUse
RemainingSpace = 
		DiskSpaceUse -> DiskPart
DynamicDiskPart DiskSpaceUse
RemainingSpace
	DynamicDiskPart (Percent Int
a) <> DiskPart
_ = DiskSpaceUse -> DiskPart
DynamicDiskPart (Int -> DiskSpaceUse
Percent Int
a)
	DiskPart
_ <> DynamicDiskPart (Percent Int
b) = DiskSpaceUse -> DiskPart
DynamicDiskPart (Int -> DiskSpaceUse
Percent Int
b)
	DynamicDiskPart DiskSpaceUse
RemainingSpace <> DiskPart
_ = DiskSpaceUse -> DiskPart
DynamicDiskPart DiskSpaceUse
RemainingSpace
	DiskPart
_ <> DynamicDiskPart DiskSpaceUse
RemainingSpace = DiskSpaceUse -> DiskPart
DynamicDiskPart DiskSpaceUse
RemainingSpace

instance Monoid DiskPart
  where
	mempty :: DiskPart
mempty = DiskPart
FixedDiskPart
	mappend :: DiskPart -> DiskPart -> DiskPart
mappend = forall a. Semigroup a => a -> a -> a
(Sem.<>)

-- | Make a partition use some percentage of the size of the disk
-- (less all fixed size partitions), or the remaining space in the disk.
useDiskSpace :: PartSpec DiskPart -> DiskSpaceUse -> PartSpec DiskPart
useDiskSpace :: PartSpec DiskPart -> DiskSpaceUse -> PartSpec DiskPart
useDiskSpace (Maybe FilePath
mp, MountOpts
o, PartSize -> Partition
p, DiskPart
_) DiskSpaceUse
diskuse = (Maybe FilePath
mp, MountOpts
o, PartSize -> Partition
p, DiskSpaceUse -> DiskPart
DynamicDiskPart DiskSpaceUse
diskuse)

-- | Default partition size when not otherwize specified is 128 MegaBytes.
defSz :: PartSize
defSz :: PartSize
defSz = ByteSize -> PartSize
MegaBytes ByteSize
128

-- | When a partition is sized to fit the files that live in it,
-- this fudge factor is added to the size of the files. This is necessary
-- since filesystems have some space overhead.
-- 
-- Add 2% for filesystem overhead. Rationalle for picking 2%:
-- A filesystem with 1% overhead might just sneak by as acceptable.
-- Double that just in case. Add an additional 3 mb to deal with
-- non-scaling overhead of filesystems (eg, superblocks). 
-- Add an additional 200 mb for temp files, journals, etc.
fudgeSz :: PartSize -> PartSize
fudgeSz :: PartSize -> PartSize
fudgeSz (MegaBytes ByteSize
n) = ByteSize -> PartSize
MegaBytes (ByteSize
n forall a. Num a => a -> a -> a
+ ByteSize
n forall a. Integral a => a -> a -> a
`div` ByteSize
100 forall a. Num a => a -> a -> a
* ByteSize
2 forall a. Num a => a -> a -> a
+ ByteSize
3 forall a. Num a => a -> a -> a
+ ByteSize
200)
fudgeSz (Bytes ByteSize
n) = PartSize -> PartSize
fudgeSz (ByteSize -> PartSize
toPartSize ByteSize
n)

alignTo :: Alignment -> PartSize -> ByteSize
alignTo :: Alignment -> PartSize -> ByteSize
alignTo Alignment
_ (Bytes ByteSize
n) = ByteSize
n -- no alignment done for Bytes
alignTo (Alignment ByteSize
alignment) PartSize
partsize
	| ByteSize
alignment forall a. Ord a => a -> a -> Bool
< ByteSize
1 = ByteSize
n
	| Bool
otherwise = case forall a. Integral a => a -> a -> a
rem ByteSize
n ByteSize
alignment of
		ByteSize
0 -> ByteSize
n
		ByteSize
r -> ByteSize
n forall a. Num a => a -> a -> a
- ByteSize
r forall a. Num a => a -> a -> a
+ ByteSize
alignment
  where
	n :: ByteSize
n = PartSize -> ByteSize
fromPartSize PartSize
partsize