Add "setup disk" command

Enables sending format and partition commands to the ogServer.

Syntax: ogcli setup disk {options}

Options are:

  --type [{dos,gpt}]    Disk partition scheme
  --num [NUM]           Disk number (defaults to 1)
  --format [FORMAT]     Indicates which partitions to reformat if they
                        are already present. Use --part alone to mean all
                        partitions.
  --part PART [PART ...]
                        Partition definition (syntax:
                        "num,part_scheme,fs,size")

	size is specified as a positive digit followed by a size unit
	like [MGT]

		512M, 50G, 1T, etc.

	ogcli accepts {LINUX,WINDOWS,EFI,CACHE} as part types and {EXT4,
	NTFS, FAT32, CACHE} as filesystem types inside the partition
	definition string

	1,LINUX,EXT4,50G => partition 1, of type linux with ext4 fs of
			    50 GB size

	If you wanted to add a partition cache (OpenGnsys usually
	mandates this partition to be number 4) you can omit partitions
	up to that number and ogcli will fill with EMPTY partitions.

	So if you define the following parts:
		--part 1,...
		--part 4,...
	ogcli will fill parts 2 and 3 with EMPTY part type and fs.

Client selection argument follow previous used options

  --center-id CENTER_ID
  --room-id ROOM_ID
  --client-ip CLIENT_IP

2 files changed, 128 insertions, 0 deletions

diff --git a/cli/cli.py b/cli/cli.py
index 1fa411f..61db374 100644
--- a/cli/cli.py
+++ b/cli/cli.py
@@ -110,3 +110,12 @@ class OgCLI():
 
 		if parsed_args.create_obj == 'image':
 			OgImage.create_image(self.rest, args[1:])
+
+	def setup(self, args):
+		choices = ['disk']
+		parser = argparse.ArgumentParser(prog='ogcli setup')
+		parser.add_argument('setup_obj', choices=choices)
+		parsed_args = parser.parse_args([args[0]])
+
+		if parsed_args.setup_obj == 'disk':
+			OgDisk.setup_disk(self.rest, args[1:])
diff --git a/cli/objects/disks.py b/cli/objects/disks.py
index b13e116..0aa5a50 100644
--- a/cli/objects/disks.py
+++ b/cli/objects/disks.py
@@ -7,6 +7,7 @@
 #
 
 import argparse
+import re
 
 class OgDisk():
 
@@ -22,3 +23,121 @@ class OgDisk():
 
 		r = rest.get('/client/setup', payload=payload)
 		print(r.text)
+
+	@staticmethod
+	def setup_disk(rest, args):
+		def parse_size(size):
+			size = size.upper()
+			units = {"M": 10**3, "G": 10**6, "T": 10**9} # Mapped to K
+			# size = re.sub(r'(\d+)([MGT])', r'\1 \2', size)
+			match = re.match(r'(\d+)([MGT])', size)
+			if match:
+				if len(match.groups()) == 2:
+					number, unit = match.groups()
+					return str(int(float(number)*units[unit]))
+			print(f'Error parsing size {size}. Aborting.')
+			return None
+
+		disk_type_map = {'dos': 'MSDOS', 'gpt': 'GPT'}
+		part_types = ['LINUX', 'EFI', 'WINDOWS', 'CACHE']
+		fs_types = ['EXT4', 'FAT32', 'NTFS', 'CACHE']
+
+		parser = argparse.ArgumentParser()
+		parser.add_argument('--type',
+				    nargs='?',
+				    required=True,
+				    choices=['dos','gpt'],
+				    help='Disk partition scheme')
+		parser.add_argument('--num',
+				    nargs='?',
+				    default=1,
+				    help='Disk number (defaults to 1)')
+		parser.add_argument('--format',
+				    nargs='?',
+				    const=True,
+				    type=lambda x: x.split(','),
+				    help='Indicates which partitions to reformat if they are already present. '\
+					 'Use --part alone to mean all partitions.')
+		parser.add_argument('--part',
+				    nargs='+',
+				    action='append',
+				    type=lambda x: x.split(','),
+				    required=True,
+				    help='Partition definition (syntax: "num,part_scheme,fs,size")')
+		group = parser.add_argument_group('clients', 'Client selection args')
+		group.add_argument('--center-id',
+				   type=int,
+				   action='append',
+				   default=[],
+				   required=False,
+				   help='Clients from given center id')
+		group.add_argument('--room-id',
+				   type=int,
+				   action='append',
+				   default=[],
+				   required=False,
+				   help='Clients from given room id')
+		group.add_argument('--client-ip',
+				   action='append',
+				   default=[],
+				   required=False,
+				   help='Specific client IP')
+
+		parsed_args = parser.parse_args(args)
+		r = rest.get('/scopes')
+		scopes = r.json()
+		ips = set()
+
+		for center in parsed_args.center_id:
+			center_scope = scope_lookup(center, 'center', scopes)
+			ips.update(ips_in_scope(center_scope))
+		for room in parsed_args.room_id:
+			room_scope = scope_lookup(room, 'room', scopes)
+			ips.update(ips_in_scope(room_scope))
+		for l in parsed_args.client_ip:
+			ips.add(l)
+		if not ips:
+			print("No clients found")
+			return None
+
+		payload = {'clients': parsed_args.client_ip, 'type': disk_type_map[parsed_args.type], 'disk': str(parsed_args.num),
+			   'cache': '0', 'cache_size': '0', 'partition_setup': []}
+		for i, p in enumerate(parsed_args.part, start=1):
+			p = p[0]
+			part_num, code, fs, size = p[0], p[1].upper(), p[2].upper(), p[3]
+
+			if code not in part_types:
+			    print(f'Specified partition type {code} is not supported. Aborting...')
+			    return
+			if fs not in fs_types:
+			    print(f'Specified filesystem {code} is not supported. Aborting...')
+			    return
+			size = parse_size(size)
+
+			for j in range(i, int(part_num)):
+				part = {'partition': str(j), 'code':'EMPTY',
+					'filesystem': 'EMPTY', 'size': '0',
+					'format': '0'}
+				payload['partition_setup'].append(part)
+
+			if parsed_args.format is True or (type(parsed_args.format) == list and part_num in parsed_args.format):
+			    do_format = '1'
+			else:
+			    do_format = '0'
+
+			if fs == 'CACHE':
+				payload['cache'] = '1'	# Assuming flag specifying if there's cache in the setup
+				payload['cache_size'] = size
+			part = {'partition': str(p[0]), 'code':code.upper(),
+				'filesystem': fs.upper(), 'size': size,
+				'format': do_format}
+			payload['partition_setup'].append(part)
+
+		# Pad with empty partitions if no 4th part was defined
+		for i in range(len(parsed_args.part)+1, 5):
+			part = {'partition': str(i), 'code':'EMPTY',
+				'filesystem': 'EMPTY', 'size': '0',
+				'format': '0'}
+			payload['partition_setup'].append(part)
+
+		rest.post('/setup', payload=payload)