2 files changed, 45 insertions, 9 deletions
diff --git a/roles/install-grub/tasks/main.yml b/roles/install-grub/tasks/main.yml
index 974f011..608d002 100644
--- a/roles/install-grub/tasks/main.yml
+++ b/roles/install-grub/tasks/main.yml
@@ -27,20 +27,24 @@
 - name: Mount EFI
   become: true
   mount:
-      src: "{{ vyos_target_drive }}p2"
+      src: "{{ vyos_target_drive }}p{{ partition_num_efi }}"
       path: "{{ vyos_install_root }}/boot/efi"
       fstype: vfat
       state: mounted
+  when: partition_num_efi is defined
 
 - name: Install GRUB in the boot sector of {{ vyos_target_drive }}
   become: true
   command: "chroot {{ vyos_install_root }} grub-install --no-floppy --target=i386-pc --root-directory=/boot {{ loop_device.stdout }} --force"
   args:
       creates: "{{ vyos_install_root }}/boot/grub/grubenv"
+  when: (vyos_parttable_type == "mbr") or
+        (vyos_parttable_type == "hybrid")
 
-- name: Install EFI GRUB to {{ vyos_target_drive }}p2
+- name: Install EFI GRUB to {{ vyos_target_drive }}p{{ partition_num_efi }}
   become: true
   command: "chroot {{ vyos_install_root }} grub-install --no-floppy --recheck --target=x86_64-efi --force-extra-removable --root-directory=/boot --efi-directory=/boot/efi --bootloader-id='VyOS' --no-uefi-secure-boot"
+  when: partition_num_efi is defined
 
 - name: Configure GRUB
   become: true
@@ -52,8 +56,8 @@
 - name: Unmount EFI
   become: true
   mount:
-      src: "{{ vyos_target_drive }}p2"
+      src: "{{ vyos_target_drive }}p{{ partition_num_efi }}"
       path: "{{ vyos_install_root }}/boot/efi"
       fstype: vfat
       state: absent
-
+  when: partition_num_efi is defined
diff --git a/roles/setup-root-partition/tasks/main.yml b/roles/setup-root-partition/tasks/main.yml
index 57ba5ae..bf183f3 100644
--- a/roles/setup-root-partition/tasks/main.yml
+++ b/roles/setup-root-partition/tasks/main.yml
@@ -1,7 +1,38 @@
-- name: Partition disk
+- name: Partition disk (MBR)
+  become: true
+  command:
+      cmd: "sfdisk -q -X dos {{ vyos_raw_img }}"
+      stdin: "3,+,L,*"
+  when: vyos_parttable_type == "mbr"
+
+- name: Partition disk (GPT)
+  become: true
+  command:
+      cmd: "sgdisk -a1 -n1:34:+256M -t1:EF00 -n2:0:0:+100% -t2:8300 {{ vyos_raw_img }}"
+  when: vyos_parttable_type == "gpt"
+
+- name: Partition disk (hybrid)
   become: true
   command:
       cmd: "sgdisk -a1 -n1:34:2047 -t1:EF02 -n2:2048:+256M -t2:EF00 -n3:0:0:+100% -t3:8300 {{ vyos_raw_img }}"
+  when: vyos_parttable_type == "hybrid"
+
+- name: Set partition numbers for next actions (MBR)
+  set_fact:
+    partition_num_root: 1
+  when: vyos_parttable_type == "mbr"
+
+- name: Set partition numbers for next actions (GPT)
+  set_fact:
+    partition_num_efi: 1
+    partition_num_root: 2
+  when: vyos_parttable_type == "gpt"
+
+- name: Set partition numbers for next actions (hybrid)
+  set_fact:
+    partition_num_efi: 2
+    partition_num_root: 3
+  when: vyos_parttable_type == "hybrid"
 
 - name: Reread partitions from image and mount to loopback
   become: true
@@ -21,20 +52,21 @@
   become: true
   filesystem:
     fstype: "vfat"
-    device: "{{ vyos_target_drive }}p2"
+    device: "{{ vyos_target_drive }}p{{ partition_num_efi }}"
     opts: "-n EFI -F 32 -s 1"
- 
+  when: partition_num_efi is defined
+
 - name: Create a fileystem on root partition
   become: true
   filesystem:
     fstype: "{{ vyos_root_fstype }}"
-    device: "{{ vyos_target_drive }}p3"
+    device: "{{ vyos_target_drive }}p{{ partition_num_root }}"
     opts: "-L persistence"
 
 - name: Mount root partition
   become: true
   mount:
     name: "{{ vyos_write_root }}"
-    src: "{{ vyos_target_drive }}p3"
+    src: "{{ vyos_target_drive }}p{{ partition_num_root }}"
     fstype: "{{ vyos_root_fstype }}"
     state: mounted