blob: 1926f177af96a73c985deb315993b66415e6e9ae [file] [log] [blame]
Dave Pigott281203e2014-09-17 23:45:02 +01001#! /usr/bin/python
2
3# Copyright 2014 Linaro Limited
Steve McIntyre663dc062014-10-20 11:11:47 +01004# Author: Dave Pigott <dave.pigott@linaro.org>
Dave Pigott281203e2014-09-17 23:45:02 +01005#
6# This program is free software; you can redistribute it and/or modify
7# it under the terms of the GNU General Public License as published by
8# the Free Software Foundation; either version 2 of the License, or
9# (at your option) any later version.
10#
11# This program is distributed in the hope that it will be useful,
12# but WITHOUT ANY WARRANTY; without even the implied warranty of
13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14# GNU General Public License for more details.
15#
16# You should have received a copy of the GNU General Public License
17# along with this program; if not, write to the Free Software
18# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
19# MA 02110-1301, USA.
20
21import psycopg2
22import psycopg2.extras
Dave Pigott9b73f3a2014-09-18 22:55:42 +010023from errors import CriticalError
Dave Pigott281203e2014-09-17 23:45:02 +010024
25class VlanDB:
26 def __init__(self, db_name="vland", username="vland"):
27 try:
Steve McIntyree38f6222014-11-27 15:09:49 +000028 self.connection = psycopg2.connect(database=db_name, user=username)
29 self.cursor = self.connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
Dave Pigott281203e2014-09-17 23:45:02 +010030 except Exception as e:
31 print "Failed to access database: %s" % e
32
33 def __del__(self):
34 self.cursor.close()
35 self.connection.close()
36
37 def create_switch(self, name, admin_password, cli_password):
Dave Pigott2649a1a2014-09-18 00:04:49 +010038 try:
Steve McIntyre028af3b2014-11-27 15:16:22 +000039 self.cursor.execute("INSERT INTO switch (name, admin_pw, cli_pw) VALUES ('%s', '%s', '%s') RETURNING switch_id" %
40 (name, admin_password, cli_password))
Dave Pigott2649a1a2014-09-18 00:04:49 +010041 switch_id = self.cursor.fetchone()[0]
42 self.connection.commit()
43 except:
44 self.connection.rollback()
45 raise
Dave Pigott281203e2014-09-17 23:45:02 +010046 return switch_id
47
Dave Pigott9b73f3a2014-09-18 22:55:42 +010048 def create_port(self, name, switch_id, is_locked, is_trunk, current_vlan_id, base_vlan_id):
Dave Pigott2649a1a2014-09-18 00:04:49 +010049 try:
Steve McIntyre028af3b2014-11-27 15:16:22 +000050 self.cursor.execute("INSERT INTO port (name, switch_id, is_locked, is_trunk, current_vlan_id, base_vlan_id)"
51 "VALUES ('%s', '%s', '%s', '%s', '%s', '%s') RETURNING port_id" %
Dave Pigott9b73f3a2014-09-18 22:55:42 +010052 (name, switch_id, is_locked, is_trunk, current_vlan_id, base_vlan_id))
Dave Pigott2649a1a2014-09-18 00:04:49 +010053 port_id = self.cursor.fetchone()[0]
54 self.connection.commit()
55 except:
56 self.connection.rollback()
57 raise
Dave Pigott281203e2014-09-17 23:45:02 +010058 return port_id
59
60 def create_vlan(self, name, tag, is_default_vlan):
Dave Pigott2649a1a2014-09-18 00:04:49 +010061 try:
Steve McIntyre028af3b2014-11-27 15:16:22 +000062 self.cursor.execute("INSERT INTO vlan (name, tag, is_default_vlan)"
63 "VALUES ('%s', '%s', '%s') RETURNING vlan_id" %
Dave Pigott2649a1a2014-09-18 00:04:49 +010064 (name, tag, is_default_vlan))
65 vlan_id = self.cursor.fetchone()[0]
66 self.connection.commit()
67 except:
68 self.connection.rollback()
69 raise
Dave Pigott281203e2014-09-17 23:45:02 +010070 return vlan_id
71
72 def _delete_row(self, table, field, value):
Dave Pigott2649a1a2014-09-18 00:04:49 +010073 try:
74 self.cursor.execute("DELETE FROM %s WHERE %s = %s" % (table, field, value))
75 self.connection.commit()
76 except:
77 self.connection.rollback()
78 raise
Dave Pigott281203e2014-09-17 23:45:02 +010079
80 def delete_switch(self, switch_id):
81 self._delete_row("switch", "switch_id", switch_id)
82
83 def delete_port(self, port_id):
84 self._delete_row("port", "port_id", port_id)
85
86 def delete_vlan(self, vlan_id):
87 self._delete_row("vlan", "vlan_id", vlan_id)
88
Dave Pigott9b73f3a2014-09-18 22:55:42 +010089 def _get_element(self, select_field, table, compare_field, value):
Dave Pigott281203e2014-09-17 23:45:02 +010090 self.cursor.execute("SELECT %s FROM %s WHERE %s = %s" % (select_field, table, compare_field, value))
91 return self.cursor.fetchone()[0]
92
93 def get_switch_id(self, name):
Dave Pigott9b73f3a2014-09-18 22:55:42 +010094 return self._get_element("switch_id", "switch", "name", name)
Dave Pigott281203e2014-09-17 23:45:02 +010095
96 def get_port_id(self, name):
Dave Pigott9b73f3a2014-09-18 22:55:42 +010097 return self._get_element("port_id", "port", "name", name)
Dave Pigott281203e2014-09-17 23:45:02 +010098
99 def get_vlan_id(self, name):
Dave Pigott9b73f3a2014-09-18 22:55:42 +0100100 return self._get_element("vlan_id", "vlan", "name", name)
Dave Pigott281203e2014-09-17 23:45:02 +0100101
102 def get_switch_name(self, switch_id):
Dave Pigott9b73f3a2014-09-18 22:55:42 +0100103 return self._get_element("name", "switch", "switch_id", switch_id)
Dave Pigott281203e2014-09-17 23:45:02 +0100104
105 def get_port_name(self, port_id):
Dave Pigott9b73f3a2014-09-18 22:55:42 +0100106 return self._get_element("port_name", "port", "port_id", port_id)
Dave Pigott281203e2014-09-17 23:45:02 +0100107
108 def get_vlan_name(self, vlan_id):
Dave Pigott9b73f3a2014-09-18 22:55:42 +0100109 return self._get_element("vlan_name", "vlan", "vlan_id", vlan_id)
110
111 def _get_row(self, table, field, value):
112 self.cursor.execute("SELECT * FROM %s WHERE %s = %s" % (table, field, value))
113 return self.cursor.fetchone()
114
115 def get_switch(self, switch_id):
116 return self._get_row("switch", "switch_id", switch_id)
117
118 def get_port(self, port_id):
119 return self._get_row("port", "port_id", port_id)
120
121 def get_vlan(self, vlan_id):
122 return self._get_row("vlan", "vlan_id", vlan_id)
123
124 def set_vlan(self, port_id, vlan_id):
125 port = self.get_port(port_id)
126 if port == None:
127 raise("Port %s does not exist" % port_id)
128
129 if port["is_trunk"] or port["is_locked"]:
130 raise CriticalError("The port is locked")
131
132 vlan = self.get_vlan(vlan_id)
133 if vlan == None:
134 raise CriticalError("VLAN %s does not exist" % vlan_id)
135
136 try:
137 self.cursor.execute("UPDATE port SET current_vlan_id=%s WHERE port_id=%s" %(vlan_id, port_id))
138 except:
139 self.connection.rollback()
140 raise
141
142 def restore_default_vlan(self, port_id):
143 port = self.get_port(port_id)
144 if port == None:
145 raise CriticalError("Port %s does not exist")
146
147 if port["is_trunk"] or port["is_locked"]:
148 raise CriticalError("The port is locked")
149
150 try:
151 self.cursor.execute("UPDATE port SET current_vlan_id=base_vlan_id WHERE port=%s" % port_id)
152 except:
153 self.connection.rollback()
154 raise
155
156
Dave Pigott281203e2014-09-17 23:45:02 +0100157
158 def _dump_table(self, table):
159 result = []
160 self.cursor.execute("SELECT * FROM %s" % table)
Dave Pigott281203e2014-09-17 23:45:02 +0100161 record = self.cursor.fetchone()
162 while record != None:
Steve McIntyree73eb122014-11-27 15:18:47 +0000163 result.append(record)
Dave Pigott281203e2014-09-17 23:45:02 +0100164 record = self.cursor.fetchone()
165 return result
166
167 def all_switches(self):
168 return self._dump_table("switch")
169
170 def all_ports(self):
171 return self._dump_table("port")
172
173 def all_vlans(self):
174 return self._dump_table("vlan")
Dave Pigott9b73f3a2014-09-18 22:55:42 +0100175