Support drives with multiple partitions (#43)

contrib/partitioner/partitioner.py

@@ -18,6 +18,7 @@ import os
 import sys
 import glob
 import re
+import subprocess
 
 usb_dev_pattern = ['sd.*']
 usb_part_pattern = ['sd.[1-9]*']
@@ -78,6 +79,11 @@ def usb_partition_table():
         table[partition] = int(usb_part_size(partition))
     return table
 
+def get_partition_type(partition):
+    command = ['blkid', '-o', 'value', '-s', 'TYPE', f'/dev/{partition}']
+    output = subprocess.run(command, stdout=subprocess.PIPE).stdout
+    return output.decode('utf-8').rstrip()
+
 '''
 Main Entrypoint
 '''
@@ -95,9 +101,10 @@ def main():
         print('Data mount exists')
 
 
-    if len(usb_devs()) == 1:
+    # STEP to check If  no partitions found
     if len(usb_partitions()) < 1:
         try:
+            # hardcoded sda for now.. may try out usb_devs() later and see for a more smoother experience
             print("Running parted...")
             os.system('/sbin/parted -s /dev/sda mkpart p ext4 3 100%');
             print('Making first partition');
@@ -106,13 +113,42 @@ def main():
             print('Error running parted');
             sys.exit(1);
     else:
-            print("Already partitioned")
-
+        print("Already partitioned.. proceed with next step")
+    # END STEP for no partitions found
     
+    # Check for more than 1 partition (and one drive only)
+    # if so get rid of it all
+    if len(usb_devs()) != 1:
+        print('More than one drive is not supported');
+        sys.exit(1);        
+    else:
+        if len(usb_partitions()) > 1:
+            print('More than one partition found! Flattening the curve..');
+            device_name=str(usb_partitions()[0])[:-1]
+            for partition in usb_partitions():
+                pn = partition.replace(device_name,'')
+                os.system('/sbin/parted -s /dev/' + device_name + ' rm ' + pn)        
 
-        if len(usb_partitions()) == 1:
+    # Check for number of partitions (only 1 partition is supported
+    if len(usb_devs()) != 1:
+        print('No drives or unexpected number of drives detected!')      
+        sys.exit(1);
+    else:
+        # Proceed to mounting the drive (and formatting for good measure :D)
         try:
             os.system('/bin/mount /dev/' + usb_partitions()[0] + ' /mnt/data')
+
+            # If not EXT4, delete partition and format (format for good measure)
+            if not get_partition_type(usb_partitions()[0]) == "ext4":
+                print('Not an EXT filesystem, remove all partitions')
+                device_name=str(usb_partitions()[0])[:-1]
+                for p in usb_partitions():
+                    pn = p.replace(device_name,'')
+                    os.system('parted -s /dev/' + device_name + ' rm ' + pn);
+                print('Running parted, and recreating partition as EXT4')
+                os.system('/sbin/parted -s /dev/' + device_name + ' mkpart p ext4 3 100%');
+                os.system('/sbin/mkfs.ext4 -F /dev/' + device_name + '1');
+
             # if .rekt exists or bitcoin directory doesnt exist (because the drive is a factory default)
             # then wipe the drive and format it with EXT4
             if os.path.exists('/mnt/data/.rekt') or not os.path.exists('/mnt/data/bitcoin'):
@@ -124,9 +160,10 @@ def main():
                 # remount
                 os.system('/bin/mount -t ext4 /dev/' + usb_partitions()[0] + ' /mnt/data')
                 '''
-                    Get Size of SDA and partition info
+                Get Size of the first partition and partition info
                 '''
-                    first_part = dev_size('sda') / (1000*1000)
+                device_name=str(usb_partitions()[0])[:-1]
+                first_part = dev_size(device_name) / (1000*1000)
                 prune_setting = int(first_part / 2)
 
                 if first_part < 512000:
@@ -158,7 +195,7 @@ def main():
 
             else:
                 '''
-                    No need to do anything
+                No need to do anything as drive is intact
                 '''
                 print('REKT file does not exist OR bitcoin folder exists so we will preserve it. But lets check for all the other folders')
 
@@ -200,11 +237,12 @@ def main():
             os.system('/bin/umount /mnt/data')
         except:
             print("Error mounting the directory")
-
+            sys.exit(1);
+        # END: Drive prep
 
         # If volume not mounted
         if not os.path.exists(' /mnt/data/lost+found'):
-            os.system('/bin/mount -t ext4 /dev/sda1 /mnt/data')
+            os.system('/bin/mount -t ext4 ' + usb_partitions()[0] + ' /mnt/data')
 
         # Get UUID of the partition we just created
         partitions = usb_partitions()
@@ -236,8 +274,6 @@ def main():
         os.system('/bin/ln -s /mnt/data/lnd ' + homedirpath + '/lnd')
         os.system('/bin/ln -s /mnt/data/tor ' + homedirpath + '/tor')
         os.system('/bin/ln -s /mnt/data/nginx ' + homedirpath + '/nginx')
-    else:
-        print('No drives or unexpected number of drives detected!')
         
 '''
 Actual entrypoint