Added computer name functions

2 files changed, 40 insertions, 5 deletions

diff --git a/machammer/network.py b/machammer/network.py
index e1a9058..dc85b2d 100644
--- a/machammer/network.py
+++ b/machammer/network.py
@@ -1,11 +1,18 @@
 # -*- coding: utf-8 -*-
 """Network-related machammer functions."""
 
+import re
 from .process import kill
 from .functions import call, check_output
 from .system_profiler import SystemProfile
 
-NETWORKSETUP = '/usr/sbin/networksetup'
+
+def networksetup(*args):
+    return check_output('/usr/sbin/networksetup', *args)
+
+
+def systemsetup(*args):
+    return check_output('/usr/sbin/systemsetup', *args)
 
 
 def get_ports(type='Ethernet'):
@@ -18,12 +25,12 @@ def set_wifi_power(on=True):
     """Set AirPort power to on (True) or off (False)"""
     state = 'on' if on else 'off'
     for i in get_ports('AirPort'):
-        call(NETWORKSETUP, '-setairportpower', i['interface'], state)
+        networksetup('-setairportpower', i['interface'], state)
 
 
 def disable_wifi(port='en1'):
-    call(NETWORKSETUP, '-setairportpower', port, 'off')
-    call(NETWORKSETUP, '-setnetworkserviceenabled', 'Wi-Fi', 'off')
+    networksetup('-setairportpower', port, 'off')
+    networksetup('-setnetworkserviceenabled', 'Wi-Fi', 'off')
 
 
 def get_wifi_power():
@@ -31,7 +38,7 @@ def get_wifi_power():
     results = []
     for i in get_ports('AirPort'):
         iface = i['interface']
-        r = check_output(NETWORKSETUP, '-getairportpower', iface)
+        r = networksetup('-getairportpower', iface)
         results.append(r.split(': ')[1])
 
     return 'On' in results
@@ -68,3 +75,23 @@ def get_primary(port=None):
 def flush_dns():
     """Flush the DNS cache."""
     kill('mDNSResponder', 'HUP')
+
+
+def get_computer_name():
+    """Return the Computer Name of this Mac"""
+    return networksetup('-getcomputername')
+
+
+def set_computer_name(name, hostname=True):
+    networksetup('-setcomputername', name)
+    if hostname:
+        set_hostname(name)
+
+
+def set_hostname(name):
+    systemsetup('-setlocalsubnetname', name)
+
+
+def fix_computer_name(name):
+    """Removes the (\d) from the computer name and local hostname"""
+    return re.sub(r'\s\(\d+\)', '', name)
diff --git a/tests.py b/tests.py
index a586ed6..081bd83 100755
--- a/tests.py
+++ b/tests.py
@@ -83,6 +83,14 @@ class SystemProfilerTestCase(TestCase):
 
 
 class NetworkTestCase(TestCase):
+    def test_get_computer_name(self):
+        name = network.get_computer_name()
+        self.assertEquals(name, 'lalalala')
+        
+    def test_fix_computer_name(self):
+        name = network.fix_computer_name('Computer (2)')
+        self.assertEqual(name, 'Computer')
+        
     def test_wifi_enable(self):
         """Turn wifi power on, check that it's on"""
         network.set_wifi_power(True)