Merge pull request #2298 from jestabro/disk-by-id

smoketest: T5607: support getting SCSI device by drive-id

6 files changed, 236 insertions, 2 deletions

diff --git a/op-mode-definitions/disks.xml.in b/op-mode-definitions/disks.xml.in
index 117ac5065..8a1e2c86f 100644
--- a/op-mode-definitions/disks.xml.in
+++ b/op-mode-definitions/disks.xml.in
@@ -5,6 +5,26 @@
       <help>Format a device</help>
     </properties>
     <children>
+      <node name="by-id">
+        <properties>
+          <help>Find disk by ending of id string</help>
+        </properties>
+        <children>
+          <tagNode name="disk">
+            <properties>
+              <help>Format a disk drive</help>
+            </properties>
+            <children>
+              <tagNode name="like">
+                <properties>
+                  <help>Format this disk the same as another disk</help>
+                </properties>
+                <command>sudo ${vyos_op_scripts_dir}/format_disk.py --by-id --target $4 --proto $6</command>
+              </tagNode>
+            </children>
+          </tagNode>
+        </children>
+      </node>
       <tagNode name="disk">
         <properties>
           <help>Format a disk drive</help>
diff --git a/op-mode-definitions/raid.xml.in b/op-mode-definitions/raid.xml.in
new file mode 100644
index 000000000..5d0c9ef3d
--- /dev/null
+++ b/op-mode-definitions/raid.xml.in
@@ -0,0 +1,69 @@
+<?xml version="1.0"?>
+<interfaceDefinition>
+  <node name="add">
+    <children>
+      <tagNode name="raid">
+        <properties>
+          <help>Add a RAID set element</help>
+          <completionHelp>
+            <script>${vyos_completion_dir}/list_raidset.sh</script>
+          </completionHelp>
+        </properties>
+        <children>
+          <node name="by-id">
+            <properties>
+              <help>Add a member by disk id to a RAID set</help>
+            </properties>
+            <children>
+              <tagNode name="member">
+                <properties>
+                  <help>Add a member to a RAID set</help>
+                </properties>
+                <command>sudo ${vyos_op_scripts_dir}/raid.py add --raid-set-name $3 --by-id --member $6</command>
+              </tagNode>
+            </children>
+          </node>
+          <tagNode name="member">
+            <properties>
+              <help>Add a member to a RAID set</help>
+            </properties>
+            <command>sudo ${vyos_op_scripts_dir}/raid.py add --raid-set-name $3 --member $5</command>
+          </tagNode>
+        </children>
+      </tagNode>
+    </children>
+  </node>
+  <node name="delete">
+    <children>
+      <tagNode name="raid">
+        <properties>
+          <help>Add a RAID set element</help>
+          <completionHelp>
+            <script>${vyos_completion_dir}/list_raidset.sh</script>
+          </completionHelp>
+        </properties>
+        <children>
+          <node name="by-id">
+            <properties>
+              <help>Add a member by disk id to a RAID set</help>
+            </properties>
+            <children>
+              <tagNode name="member">
+                <properties>
+                  <help>Add a member to a RAID set</help>
+                </properties>
+                <command>sudo ${vyos_op_scripts_dir}/raid.py delete --raid-set-name $3 --by-id --member $6</command>
+              </tagNode>
+            </children>
+          </node>
+          <tagNode name="member">
+            <properties>
+              <help>Add a member to a RAID set</help>
+            </properties>
+            <command>sudo ${vyos_op_scripts_dir}/raid.py delete --raid-set-name $3 --member $5</command>
+          </tagNode>
+        </children>
+      </tagNode>
+    </children>
+  </node>
+</interfaceDefinition>
diff --git a/python/vyos/raid.py b/python/vyos/raid.py
new file mode 100644
index 000000000..7fb794817
--- /dev/null
+++ b/python/vyos/raid.py
@@ -0,0 +1,71 @@
+# Copyright 2023 VyOS maintainers and contributors <maintainers@vyos.io>
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this library.  If not, see <http://www.gnu.org/licenses/>.
+
+from vyos.utils.disk import device_from_id
+from vyos.utils.process import cmd
+
+def raid_sets():
+    """
+    Returns a list of RAID sets
+    """
+    with open('/proc/mdstat') as f:
+        return [line.split()[0].rstrip(':') for line in f if line.startswith('md')]
+
+def raid_set_members(raid_set_name: str):
+    """
+    Returns a list of members of a RAID set
+    """
+    with open('/proc/mdstat') as f:
+        for line in f:
+            if line.startswith(raid_set_name):
+                return [l.split('[')[0] for l in line.split()[4:]]
+    return []
+
+def partitions():
+    """
+    Returns a list of partitions
+    """
+    with open('/proc/partitions') as f:
+        p = [l.strip().split()[-1] for l in list(f) if l.strip()]
+    p.remove('name')
+    return p
+
+def add_raid_member(raid_set_name: str, member: str, by_id: bool = False):
+    """
+    Add a member to an existing RAID set
+    """
+    if by_id:
+        member = device_from_id(member)
+    if raid_set_name not in raid_sets():
+        raise ValueError(f"RAID set {raid_set_name} does not exist")
+    if member not in partitions():
+        raise ValueError(f"Partition {member} does not exist")
+    if member in raid_set_members(raid_set_name):
+        raise ValueError(f"Partition {member} is already a member of RAID set {raid_set_name}")
+    cmd(f'mdadm --add /dev/{raid_set_name} /dev/{member}')
+    disk = cmd(f'lsblk -ndo PKNAME /dev/{member}')
+    cmd(f'grub-install /dev/{disk}')
+
+def delete_raid_member(raid_set_name: str, member: str, by_id: bool = False):
+    """
+    Delete a member from an existing RAID set
+    """
+    if by_id:
+        member = device_from_id(member)
+    if raid_set_name not in raid_sets():
+        raise ValueError(f"RAID set {raid_set_name} does not exist")
+    if member not in raid_set_members(raid_set_name):
+        raise ValueError(f"Partition {member} is not a member of RAID set {raid_set_name}")
+    cmd(f'mdadm --remove /dev/{raid_set_name} /dev/{member}')
diff --git a/python/vyos/utils/disk.py b/python/vyos/utils/disk.py
new file mode 100644
index 000000000..ee540b107
--- /dev/null
+++ b/python/vyos/utils/disk.py
@@ -0,0 +1,23 @@
+# Copyright 2023 VyOS maintainers and contributors <maintainers@vyos.io>
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# This library is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public
+# License along with this library.  If not, see <http://www.gnu.org/licenses/>.
+
+from pathlib import Path
+
+def device_from_id(id):
+    """ Return the device name from (partial) disk id """
+    path = Path('/dev/disk/by-id')
+    for device in path.iterdir():
+        if device.name.endswith(id):
+            return device.readlink().stem
diff --git a/src/op_mode/format_disk.py b/src/op_mode/format_disk.py
index 31ceb196a..dc3c96322 100755
--- a/src/op_mode/format_disk.py
+++ b/src/op_mode/format_disk.py
@@ -24,6 +24,7 @@ from vyos.utils.io import ask_yes_no
 from vyos.utils.process import call
 from vyos.utils.process import cmd
 from vyos.utils.process import DEVNULL
+from vyos.utils.disk import device_from_id
 
 def list_disks():
     disks = set()
@@ -77,12 +78,18 @@ if __name__ == '__main__':
     group = parser.add_argument_group()
     group.add_argument('-t', '--target', type=str, required=True, help='Target device to format')
     group.add_argument('-p', '--proto', type=str, required=True, help='Prototype device to use as reference')
+    parser.add_argument('--by-id', action='store_true', help='Specify device by disk id')
     args = parser.parse_args()
+    target = args.target
+    proto = args.proto
+    if args.by_id:
+        target = device_from_id(target)
+        proto = device_from_id(proto)
 
-    target_disk = args.target
+    target_disk = target
     eligible_target_disks = list_disks()
 
-    proto_disk = args.proto
+    proto_disk = proto
     eligible_proto_disks = eligible_target_disks.copy()
     eligible_proto_disks.remove(target_disk)
 
diff --git a/src/op_mode/raid.py b/src/op_mode/raid.py
new file mode 100755
index 000000000..fed8ae2c3
--- /dev/null
+++ b/src/op_mode/raid.py
@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+#
+# Copyright (C) 2023 VyOS maintainers and contributors
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 or later as
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+#
+import sys
+
+import vyos.opmode
+from vyos.raid import add_raid_member
+from vyos.raid import delete_raid_member
+
+def add(raid_set_name: str, member: str, by_id: bool = False):
+    try:
+        add_raid_member(raid_set_name, member, by_id)
+    except ValueError as e:
+        raise vyos.opmode.IncorrectValue(str(e))
+
+def delete(raid_set_name: str, member: str, by_id: bool = False):
+    try:
+        delete_raid_member(raid_set_name, member, by_id)
+    except ValueError as e:
+        raise vyos.opmode.IncorrectValue(str(e))
+
+if __name__ == '__main__':
+    try:
+        res = vyos.opmode.run(sys.modules[__name__])
+        if res:
+            print(res)
+    except (ValueError, vyos.opmode.Error) as e:
+        print(e)
+        sys.exit(1)
+