CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
# -*- coding: binary -*-
2
require 'singleton'
3
require 'rex'
4

5
module Rex
6

7
###
8
#
9
# This class manages service allocation and interaction.  This class can be
10
# used to start HTTP servers and manage them and all that stuff.  Yup.
11
#
12
###
13
class ServiceManager < Hash
14

15
  #
16
  # This class is a singleton.
17
  #
18
  include Singleton
19

20
  #
21
  # Calls the instance method to start a service.
22
  #
23
  def self.start(klass, *args, **kwargs)
24
    self.instance.start(klass, *args, **kwargs)
25
  end
26

27
  #
28
  # Calls the instance method to stop a service.
29
  #
30
  def self.stop(klass, *args, **kwargs)
31
    self.instance.stop(klass, *args, **kwargs)
32
  end
33

34
  #
35
  # Stop a service using the alias that's associated with it.
36
  #
37
  def self.stop_by_alias(als)
38
    self.instance.stop_by_alias(als)
39
  end
40

41
  #
42
  # Stop the supplied service instance.
43
  #
44
  def self.stop_service(service)
45
    self.instance.stop_service(service)
46
  end
47

48
  #
49
  # Starts a service and assigns it a unique name in the service hash.
50
  #
51
  def start(klass, *args, **kwargs)
52
    # Get the hardcore alias.
53
    hals = "#{klass}" + klass.hardcore_alias(*args, **kwargs)
54

55
    # Has a service already been constructed for this guy?  If so, increment
56
    # its reference count like it aint no thang.
57
    if (inst = self[hals])
58
      inst.ref
59
      return inst
60
    end
61

62
    inst = klass.new(*args, **kwargs)
63
    als  = inst.alias
64

65
    # Find an alias that isn't taken.
66
    if (self[als])
67
      cnt  = 1
68
      cnt += 1 while (self[als + " #{cnt}"])
69
      als  = inst.alias + " #{cnt}"
70
    end
71

72
    # Extend the instance as a service.
73
    inst.extend(Rex::Service)
74

75
    # Re-aliases the instance.
76
    inst.alias = als
77

78
    # Fire up the engines.  If an error occurs an exception will be
79
    # raised.
80
    inst.start
81

82
    # Alias associate and initialize reference counting
83
    self[als] = self[hals] = inst.refinit
84

85
    # Pass the caller a reference
86
    inst.ref
87

88
    inst
89
  end
90

91
  #
92
  # Stop a service using a given klass and arguments.  These should mirror
93
  # what was originally passed to start exactly.  If the reference count of
94
  # the service drops to zero the service will be destroyed.
95
  #
96
  def stop(klass, *args, **kwargs)
97
    stop_service(hals[hardcore_alias(klass, *args, **kwargs)])
98
  end
99

100
  #
101
  # Stops a service using the provided alias.
102
  #
103
  def stop_by_alias(als)
104
    stop_service(self[als])
105
  end
106

107
  #
108
  # Stops a service instance.
109
  #
110
  def stop_service(inst)
111
    # Stop the service and be done wif it, but only if the number of
112
    # references has dropped to zero
113
    if (inst)
114
      # Since the instance may have multiple aliases, scan through
115
      # all the pairs for matching stuff.
116
      self.each_pair { |cals, cinst|
117
        self.delete(cals) if (inst == cinst)
118
      }
119

120
      # Lose the list-held reference to the instance
121
      inst.deref
122

123
      return true
124
    end
125

126
    # Return false if the service isn't there
127
    return false
128
  end
129

130
  #
131
  # Overrides the builtin 'each' operator to avoid the following exception on Ruby 1.9.2+
132
  #    "can't add a new key into hash during iteration"
133
  #
134
  def each(&block)
135
    list = []
136
    self.keys.sort.each do |sidx|
137
      list << [sidx, self[sidx]]
138
    end
139
    list.each(&block)
140
  end
141

142
protected
143

144
  #
145
  # Returns the alias for a given service instance.
146
  #
147
  def hardcore_alias(klass, *args)
148
    "__#{klass.name}#{args}"
149
  end
150

151
end
152

153
end
154

155