1from pdu import PDU, PDUPort, PDUState 

2from easysnmp import snmp_get, snmp_set, snmp_walk 

3 

4 

5def _is_int(s): 

6 try: 

7 s = int(s) 

8 return True 

9 except ValueError: 

10 return False 

11 

12 

13class BaseSnmpPDU(PDU): 

14 def __init__(self, name, hostname, oid_outlets_label_base, 

15 community="private"): 

16 self.hostname = hostname 

17 self.oid_outlets_label_base = oid_outlets_label_base 

18 self.community = community 

19 

20 super().__init__(name) 

21 

22 @property 

23 def ports(self): 

24 ports = [] 

25 

26 try: 

27 names = [x.value for x in snmp_walk(self.oid_outlets_label_base, 

28 hostname=self.hostname, 

29 community=self.community, 

30 version=1)] 

31 except SystemError as e: 

32 raise ValueError(f"The snmp_walk() call failed with the following error: {e}") 

33 

34 for i, name in enumerate(names): 

35 ports.append(PDUPort(self, i+1, name)) 

36 

37 return ports 

38 

39 def _port_spec_to_int(self, port_spec): 

40 if _is_int(port_spec): 

41 return port_spec 

42 else: 

43 for port in self.ports: 

44 if port.label == port_spec: 

45 return port.port_id 

46 raise ValueError( 

47 f"{port_spec} can not be interpreted as a valid port") 

48 

49 def port_oid(self, port_id): 

50 raise ValueError("This needs to be implemented in a child class") 

51 

52 @property 

53 def action_to_snmp_value(self): 

54 return { 

55 "ON": 1, 

56 "OFF": 2, 

57 "REBOOT": 3 

58 } 

59 

60 def state_to_raw_value(self, state): 

61 value = self.action_to_snmp_value.get(state.name, None) 

62 if value is None: 

63 raise ValueError(f"The state '{state.name}' does not have an SNMP mapping") 

64 return value 

65 

66 def raw_value_to_state(self, value): 

67 val_to_state = dict([(val, key) for key, val in self.action_to_snmp_value.items()]) 

68 return getattr(PDUState, val_to_state.get(value)) 

69 

70 def set_port_state(self, port_spec, state): 

71 SNMP_INTEGER_TYPE = 'i' 

72 

73 port_id = self._port_spec_to_int(port_spec) 

74 try: 

75 return snmp_set(self.port_oid(port_id), 

76 self.state_to_raw_value(state), 

77 SNMP_INTEGER_TYPE, 

78 hostname=self.hostname, 

79 version=1, 

80 community=self.community) 

81 except SystemError as e: 

82 raise ValueError(f"The snmp_set() call failed with the following error: {e}") 

83 

84 def get_port_state(self, port_spec): 

85 port_id = self._port_spec_to_int(port_spec) 

86 try: 

87 vs = snmp_get(self.port_oid(port_id), 

88 hostname=self.hostname, 

89 version=1, 

90 community=self.community) 

91 except SystemError as e: 

92 raise ValueError(f"The snmp_get() call failed with the following error: {e}") 

93 return self.raw_value_to_state(int(vs.value)) 

94 

95 

96class SnmpPDU(BaseSnmpPDU): 

97 def __init__(self, name, config): 

98 hostname = config.get('hostname') 

99 if hostname is None: 

100 raise ValueError("Config: Missing the 'hostname' parameter") 

101 

102 oid_outlets_label_base = config.get('oid_outlets_label_base') 

103 if oid_outlets_label_base is None: 

104 raise ValueError("Config: Missing the 'oid_outlets_label_base' parameter") 

105 

106 self.oid_outlets_base = config.get('oid_outlets_base') 

107 if self.oid_outlets_base is None: 

108 raise ValueError("Config: Missing the 'oid_outlets_base' parameter") 

109 

110 community = config.get('community', 'private') 

111 

112 self.action_mapping = config.get('action_to_snmp_value', None) 

113 if self.action_mapping is not None: 

114 for action in ["ON", "OFF", "REBOOT"]: 

115 if action not in self.action_mapping: 

116 raise ValueError(f"The action '{action}' is missing from the action mapping") 

117 value = self.action_mapping[action] 

118 try: 

119 value = int(value) 

120 except Exception as e: 

121 print(e) 

122 raise ValueError(f"The value for action '{action}' should be an integer") 

123 

124 super().__init__(name, hostname, oid_outlets_label_base, community) 

125 

126 def port_oid(self, port_id): 

127 return f"{self.oid_outlets_base}.{port_id}" 

128 

129 @property 

130 def action_to_snmp_value(self): 

131 if self.action_mapping is not None: 

132 return self.action_mapping 

133 else: 

134 return super().action_to_snmp_value